[tml] Inkling model support

megatron/core/dist_checkpointing/strategies/torch.py

@@ -109,6 +109,20 @@ def register_default_torch_strategies():
 
 logger = getLogger(__name__)
 
+_COORDINATION_PROCESS_GROUP = None
+
+
+def _get_coordination_process_group():
+    """Gloo group for checkpoint plan/metadata object collectives. On an
+    NCCL-only default group these would be staged through GPU and leave
+    permanent per-peer NCCL channel buffers on the coordinator rank."""
+    global _COORDINATION_PROCESS_GROUP
+    if not torch.distributed.is_initialized():
+        return None
+    if _COORDINATION_PROCESS_GROUP is None:
+        _COORDINATION_PROCESS_GROUP = torch.distributed.new_group(backend="gloo")
+    return _COORDINATION_PROCESS_GROUP
+
 
 def flatten_state_dict(
     state_dict: ShardedStateDict,
@@ -701,7 +715,7 @@ def async_save(
         ) = save_state_dict_async_plan(
             pyt_state_dict,
             writer,
-            None,
+            _get_coordination_process_group(),
             coordinator,
             planner=MCoreSavePlanner(
                 dedup_replicated_tensors=not self.keep_only_main_replica, flatten_state_dict=False
@@ -805,6 +819,7 @@ def load(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path) -> St
         checkpoint.load_state_dict(
             pyt_state_dict,
             fsr,
+            process_group=_get_coordination_process_group(),
             planner=MCoreLoadPlanner(
                 shapes_validation_sharded_tensors=flexible_shape_sharded_tensors,
                 allow_shape_mismatch_sharded_tensors=allow_shape_mismatch_sharded_tensors,

megatron/core/optimizer/distrib_optimizer.py

@@ -537,6 +537,7 @@ def __init__(
         )
 
         self._state_offloader: Optional[OptimizerStateOffloader] = None
+        self._nvme_state_store = None
 
         # when freezing sub-models we have no real optimizer
         # but still need a stub DistributedOptimizer class
@@ -630,6 +631,15 @@ def __init__(
         if self.config.offload_optimizer_states:
             self._state_offloader = OptimizerStateOffloader(self)
 
+        if self.config.optimizer_state_nvme_dir is not None:
+            from megatron.core.optimizer.nvme_state_store import NVMeOptimizerStateStore
+
+            self._nvme_state_store = NVMeOptimizerStateStore(
+                self,
+                self.config.optimizer_state_nvme_dir,
+                self.config.optimizer_state_nvme_chunk_mb,
+            )
+
     def _get_model_param_range_map(self, param: torch.nn.Parameter):
         """
         Given a model param, get the index sub-range of the param that this
@@ -656,6 +666,8 @@ def state_dict(self):
         optimizer state (e.g., exp_avg, exp_avg_sq) are stored in a separate
         checkpoint file by calling 'save_parameter_state()'.
         """
+        if self._nvme_state_store is not None:
+            return {"nvme_state_store": True}
         inner_state_dict = self.optimizer.state_dict()
         state_dict = {}
 
@@ -741,6 +753,9 @@ def load_state_dict(self, state_dict):
         - state_order : The index of a parameter within the shared parameter
             list.
         """
+        if self._nvme_state_store is not None:
+            return
+
         if self.ddp_config.use_megatron_fsdp:
             if "param_to_group_meta" in state_dict:
                 state_dict["param_groups"] = self._param2group_meta_to_param_groups(
@@ -1247,6 +1262,8 @@ def sharded_state_dict(
 
         Regular state dict parameters are saved on DP rank 0 and loaded on all ranks.
         """
+        if self._nvme_state_store is not None:
+            return {}
         if sharding_type is not None:
             log_single_rank(
                 logger,
@@ -2486,29 +2503,29 @@ def _copy_main_params_to_model_params(self):
         # Utility method for copying group params.
         def copy_group_params(shard_main_groups, model_groups):
             for shard_main_group, model_group in zip(shard_main_groups, model_groups):
-                for shard_main_param, model_param in zip(shard_main_group, model_group):
+                self._copy_main_params_to_model_params_for(zip(shard_main_group, model_group))
 
-                    param_range_map = self._get_model_param_range_map(model_param)
-                    world_range = param_range_map["gbuf_world_in_bucket"]
+        # Copy shard groups to model groups.
+        copy_group_params(self.shard_fp32_from_float16_groups, self.model_float16_groups)
+        copy_group_params(self.shard_fp32_groups, self.model_fp32_groups)
 
-                    assert world_range.size == shard_main_param.nelement()
+    def _copy_main_params_to_model_params_for(self, pairs):
+        """Copy (shard_main_param, model_param) pairs into the param buffer."""
+        for shard_main_param, model_param in pairs:
+            param_range_map = self._get_model_param_range_map(model_param)
+            world_range = param_range_map["gbuf_world_in_bucket"]
 
-                    gbuf_index, _, bucket_id = self.model_param_gbuf_map[model_param]
-                    model_param_buffer = self.buffers[gbuf_index].buckets[bucket_id].param_data
+            assert world_range.size == shard_main_param.nelement()
 
-                    shard_model_param = model_param_buffer.view(-1)[
-                        world_range.start : world_range.end
-                    ]
+            gbuf_index, _, bucket_id = self.model_param_gbuf_map[model_param]
+            model_param_buffer = self.buffers[gbuf_index].buckets[bucket_id].param_data
 
-                    if is_float8tensor(model_param):
-                        # FP8 params are quantized in the above "quantize_param_shard" function.
-                        continue
-                    else:
-                        shard_model_param.data.copy_(shard_main_param)
+            shard_model_param = model_param_buffer.view(-1)[world_range.start : world_range.end]
 
-        # Copy shard groups to model groups.
-        copy_group_params(self.shard_fp32_from_float16_groups, self.model_float16_groups)
-        copy_group_params(self.shard_fp32_groups, self.model_fp32_groups)
+            if is_float8tensor(model_param):
+                # FP8 params are quantized in the above "quantize_param_shard" function.
+                continue
+            shard_model_param.data.copy_(shard_main_param)
 
     def _copy_main_params_to_param_buffer(self):
         """
@@ -2571,6 +2588,15 @@ def _build_model_param_to_state_dict_param_map(self, state_dict):
         return model_param_to_state_dict_param_map
 
     def _copy_model_params_to_main_params(self, state_dict=None):
+        if self._nvme_state_store is not None:
+            self._nvme_state_store.refresh_main_from_model_params(
+                lambda: self._copy_model_params_to_main_params_impl(state_dict)
+            )
+            return
+        self._copy_model_params_to_main_params_impl(state_dict)
+
+    @torch.no_grad()
+    def _copy_model_params_to_main_params_impl(self, state_dict=None):
         """
         Copy model params to main params.
 
@@ -2634,7 +2660,10 @@ def step_with_ready_grads(self) -> bool:
         """
         if self._state_offloader is not None:
             self._state_offloader.sync_before_step()
-        update_successful = super().step_with_ready_grads()
+        if self._nvme_state_store is not None:
+            update_successful = self._nvme_state_store.step()
+        else:
+            update_successful = super().step_with_ready_grads()
 
         timers = self.config.timers
         if timers is not None: