feat(optimizer): NVMe streaming of DistributedOptimizer state

megatron/training/checkpointing.py

@@ -468,6 +468,17 @@ def save_grads(save_dir, state_dict, iteration, grad_label):
                  f"from iteration {iteration:7d}")
 
 
+def _iter_nvme_state_stores(optimizer):
+    for opt in getattr(optimizer, "chained_optimizers", None) or [optimizer]:
+        store = getattr(opt, "_nvme_state_store", None)
+        if store is not None:
+            yield store
+
+
+def _nvme_state_base_dir(checkpoint_name):
+    return os.path.join(checkpoint_name, "nvme_opt_state")
+
+
 def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floating_point_operations_so_far,
                     checkpointing_context=None, pipeline_rank=None, expert_rank=None, tensor_rank=None, pipeline_parallel=None, expert_parallel=None, non_persistent_ckpt=False,
                     train_data_iterator=None, preprocess_common_state_dict_fn = None, release=False, tp_group: Optional[torch.distributed.ProcessGroup] = None, pp_group: Optional[torch.distributed.ProcessGroup] = None, dp_cp_group: Optional[torch.distributed.ProcessGroup] = None):
@@ -570,6 +581,10 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floati
         if not optimizer.is_stub_optimizer:
             optimizer.save_state_dict_to_file(optim_checkpoint_name)
 
+    if not args.no_save_optim and optimizer is not None:
+        for store in _iter_nvme_state_stores(optimizer):
+            store.save_to(_nvme_state_base_dir(checkpoint_name))
+
     async_save_request = None
     if args.async_save:
         if ckpt_type == CheckpointType.LEGACY:
@@ -1829,6 +1844,10 @@ def load_model_state_dict(module, state_dict, strict: bool):
             else:
                 optimizer.reload_model_params()
 
+    if optimizer is not None and not release and not args.finetune and not args.no_load_optim:
+        for store in _iter_nvme_state_stores(optimizer):
+            store.load_from(_nvme_state_base_dir(checkpoint_name))
+
     # rerun state
     if not ignore_rerun_state:
         try: