[Feature] Hide 75% of the communication in tensor parallelism using DoMiNo

examples/domino/domino_config.yaml

@@ -0,0 +1,105 @@
+checkpoints:
+  checkpoint_interval: 10000
+  checkpoints_path: checkpoints
+  checkpoints_path_is_shared_file_system: false
+  resume_checkpoint_path: null
+  load_lr_scheduler: false
+  load_optimizer: false
+  save_final_state: true
+  save_initial_state: false
+data_stages:
+- data:
+    dataset:
+      dataset_folder:
+      - /fsx/elie_bakouch/data/fw-edu-dedup
+    num_loading_workers: 0
+    seed: 8
+  name: stable phase
+  start_training_step: 1
+general:
+  benchmark_csv_path: null
+  consumed_train_samples: null
+  ignore_sanity_checks: true
+  project: nanotron_domino
+  run: domino_config
+  seed: 6
+  step: null
+logging:
+  iteration_step_info_interval: 1
+  log_level: info
+  log_level_replica: info
+model:
+  ddp_bucket_cap_mb: 25
+  dtype: bfloat16
+  init_method:
+    std: 0.041666666666666664
+  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: 4096
+    num_attention_heads: 32
+    num_hidden_layers: 15
+    num_key_value_heads: 8
+    pad_token_id: null
+    pretraining_tp: 2
+    rms_norm_eps: 1.0e-05
+    rope_interleaved: false
+    rope_scaling:
+      factor: 32.0
+      high_freq_factor: 4.0
+      low_freq_factor: 1.0
+      original_max_position_embeddings: 4096
+      rope_type: llama3
+    rope_theta: 500000.0
+    tie_word_embeddings: true
+    use_cache: true
+    vocab_size: 128256
+optimizer:
+  accumulate_grad_in_fp32: true
+  clip_grad: 1.0
+  learning_rate_scheduler:
+    learning_rate: 0.00005
+    lr_decay_starting_step: 50000
+    lr_decay_steps: 10000
+    lr_decay_style: linear
+    lr_warmup_steps: 1000
+    lr_warmup_style: linear
+    min_decay_lr: 0
+  optimizer_factory:
+    adam_beta1: 0.9
+    adam_beta2: 0.95
+    adam_eps: 1.0e-08
+    name: adamW
+    torch_adam_is_fused: true
+  weight_decay: 0.01
+  zero_stage: 1
+parallelism:
+  dp: 1
+  expert_parallel_size: 1
+  pp: 1
+  pp_engine: 1f1b
+  recompute_layer: false
+  tp: 8
+  tp_linear_async_communication: false
+  tp_mode: ALL_REDUCE
+  tp_recompute_allgather: false
+  domino:
+    num_input_batches: 2
+tokenizer:
+  tokenizer_max_length: null
+  tokenizer_name_or_path: meta-llama/Llama-3.2-3B
+  tokenizer_revision: null
+tokens:
+  batch_accumulation_per_replica: 2
+  limit_test_batches: 0
+  limit_val_batches: 0
+  micro_batch_size: 8
+  sequence_length: 4096
+  train_steps: 15000
+  val_check_interval: -1

src/nanotron/config/parallelism_config.py

@@ -11,6 +11,23 @@
 from nanotron.parallel.tensor_parallel.nn import TensorParallelLinearMode
 
 
+@dataclass
+class DominoArgs:
+    """
+    Domino: Eliminating Communication in LLM Training via Generic Tensor Slicing and Overlapping
+    https://arxiv.org/abs/2409.15241
+    """
+
+    # NOTE: if the number of input batches is 1,
+    # it's equivalent to non-domino mode
+    # so if you want to enable domino mode, set this to > 1
+    num_input_batches: int
+
+    def __post_init__(self):
+        assert self.num_input_batches > 1, "In order to enable domino mode, set num_input_batches > 1"
+        assert self.num_input_batches == 2, "Currently parallelism only supports 2 batches for Domino"
+
+
 @dataclass
 class ParallelismArgs:
     """Arguments related to TP/PP/DP
@@ -37,6 +54,7 @@ class ParallelismArgs:
     tp_recompute_allgather: bool = True
 
     expert_parallel_size: int = 1
+    domino: Optional[DominoArgs] = None
 
     def __post_init__(self):
         # Conservative defaults
@@ -51,3 +69,19 @@ def __post_init__(self):
             self.pp_engine = cast_str_to_pipeline_engine(self.pp_engine)
         if isinstance(self.tp_mode, str):
             self.tp_mode = TensorParallelLinearMode[self.tp_mode.upper()]
+
+        if self.is_domino_enabled is True:
+            assert self.tp > 1, "Domino requires TP > 1"
+            # NOTE: For DoMiNo since we overlapping the communication
+            # so it doesnt matter whether it's all_reduce or reduce_scatter
+            # so we just support and tested with all_reduce up to now
+            # but in principle, it should work with reduce_scatter as well
+            assert (
+                self.tp_linear_async_communication is False
+            ), "Domino requires TP linear async communication to be False"
+            # TODO: support REDUCE_SCATTER mode for Domino
+            assert self.tp_mode == TensorParallelLinearMode.ALL_REDUCE, "Domino requires TP mode to be ALL_REDUCE"
+
+    @property
+    def is_domino_enabled(self) -> bool:
+        return True if self.domino else False

src/nanotron/constants.py

@@ -10,3 +10,7 @@
 
 CHECKPOINT_FILE_NAME = "checkpoint_metadata.json"
 MODEL_CONFIG_FILE_NAME = "model_config.json"
+
+
+### FOR COMMUNICATION ###
+CUDA_STREAM_COMM_NAME = "comm_stream_{}"

src/nanotron/helpers.py

@@ -482,7 +482,9 @@ def get_profiler(config: Config):
             on_trace_ready=on_trace_ready,
             # record_shapes=True,
             # profile_memory=True,
+            with_flops=True,
             with_stack=True,
+            with_modules=True,
         )
     else:
         prof = contextlib.nullcontext()

src/nanotron/models/base.py

@@ -71,7 +71,7 @@ def get_embeddings_lm_head_tied_names(self) -> list[str]:
         Example for GPT2 model: ["model.token_position_embeddings.pp_block.token_embedding.weight", "model.lm_head.pp_block.weight"]
         """
         return []
-    
+
     def get_named_params_without_weight_decay(self) -> List[str]:
         """Return a list of named parameters that should not have weight decay applied to them."""
         return []