multi-node pp hang when enable gradient accumulation

[yuuxiaooqingg]

I tested llama3 continue training with multi-machine tp4 pp2 dp2. If I enabled grad accum operation, the training would hang. The experimental environment is: 16H800 torch 2.1.2+cu121.

checkpoints: checkpoint_interval: 200 checkpoints_path: exps/llama3_ct/ckpts checkpoints_path_is_shared_file_system: true resume_checkpoint_path: pretrained_models/Meta-Llama-3-8B save_initial_state: false no_load_optim: true

data_stages:

• name: Stable Training Stage start_training_step: 1 data: dataset: dataset_overwrite_cache: false dataset_processing_num_proc_per_process: 4 hf_dataset_config_name: null hf_dataset_or_datasets: data/cosmopedia_8193/train domain_names:
  • auto_math_text
  • khanacademy num_loading_workers: 1 seed: 42 general: benchmark_csv_path: null consumed_train_samples: 256000 ignore_sanity_checks: true project: llama_ct run: train_cosmopedia seed: 42 step: 500 logging: iteration_step_info_interval: 1 log_level: info log_level_replica: info model: ddp_bucket_cap_mb: 120 dtype: bfloat16 init_method: std: 0.025 make_vocab_size_divisible_by: 1 model_config: bos_token_id: 128000 eos_token_id: 128001 hidden_act: silu hidden_size: 4096 initializer_range: 0.02 intermediate_size: 14336 is_llama_config: true max_position_embeddings: 8192 num_attention_heads: 32 num_hidden_layers: 32 num_key_value_heads: 8 pad_token_id: null pretraining_tp: 1 rms_norm_eps: 1.0e-05 rope_scaling: null tie_word_embeddings: false use_cache: true vocab_size: 128256 optimizer: accumulate_grad_in_fp32: true optimizer_factory: adam_beta1: 0.9 adam_beta2: 0.95 adam_eps: 1.0e-08 torch_adam_is_fused: true clip_grad: 1.0 learning_rate_scheduler: learning_rate: 6.0e-05 lr_decay_steps: 400 lr_decay_style: cosine lr_warmup_steps: 100 lr_warmup_style: linear min_decay_lr: 6.0e-06 weight_decay: 0.01 zero_stage: 0 parallelism: dp: 2 pp: 2 pp_engine: 1f1b tp: 4 tp_linear_async_communication: true tp_mode: REDUCE_SCATTER profiler: null tokenizer: tokenizer_max_length: null tokenizer_name_or_path: meta-llama/Meta-Llama-3-8B tokenizer_revision: null tokens: batch_accumulation_per_replica: 2 limit_test_batches: 0 limit_val_batches: 0 micro_batch_size: 1 sequence_length: 8192 train_steps: 500 val_check_interval: -1

[yuuxiaooqingg]

@xrsrke Do you know what might be causing this?

[xrsrke]

@yuuxiaooqingg Hello. At which step did it hang?

[Pclanglais]

Same issue here. It hangs just before starting. Seems like gradient communication is not going well…

(tested on 48x4 h100s)

[Pclanglais]

Not sure if it applies in this case but I've found a fix: disable zero stage entirely. Zero stage 1 works somewhat with replicas in a smalll number of nodes setting (not 2 nor 3). Stops working as well with large distributed training sets.