align the deepseek v3 training parameters (#7)

Author: wenxie-amd

examples/deepseek_v3/exp_pretrain.yaml

@@ -16,11 +16,11 @@ modules:
     overrides:
       # log
       wandb_project: "Primus_DeepSeekV3_Pretrain"
-      disable_wandb: false
+      # disable_wandb: false
       stderr_sink_level: DEBUG
 
       # debug
-      num_layers: 4
+      # num_layers: 4
 
       # hyber parameters
       train_iters: 10
@@ -33,6 +33,12 @@ modules:
       lr_warmup_iters: 2
       lr_decay_iters: null
       lr_decay_style: cosine
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      eod_mask_loss: true
+      init_method_std: 0.008
+      norm_epsilon: 1.0e-6
 
       # parallel
       tensor_model_parallel_size: 1
@@ -41,8 +47,8 @@ modules:
 
       # data
       train_data_path: /home/azureuser/tas-public/data/deepseek-datasets/mmap_deepseekv2_datasets_text_document
-      valid_data_path: /home/azureuser/tas-public/data/deepseek-datasets/mmap_deepseekv2_datasets_text_document
-      test_data_path: /home/azureuser/tas-public/data/deepseek-datasets/mmap_deepseekv2_datasets_text_document
+      valid_data_path: null
+      test_data_path: null
 
       # fusion
       # 20250317: need latest apex in docker image
@@ -61,3 +67,4 @@ modules:
       no_save_optim: null
       no_save_rng: null
       disable_last_saving: true
+      ckpt_format: torch

primus/configs/modules/megatron/trainer_base.yaml

@@ -145,9 +145,7 @@ ddp_bucket_size: null # int
 ddp_pad_buckets_for_high_nccl_busbw: false
 ddp_average_in_collective: false
 overlap_grad_reduce: false
-delay_grad_reduce: true
 overlap_param_gather: false
-delay_param_gather: false
 overlap_param_gather_with_optimizer_step: false
 align_param_gather: true
 scatter_gather_tensors_in_pipeline: true

primus/modules/trainer/base_trainer.py

@@ -11,17 +11,14 @@
 
 
 class BaseTrainer(ABC):
-    # def get_batch_func(self):
     @abstractmethod
     def get_batch(self, data_iterator):
-        raise NotImplementedError
+        pass
 
-    # def get_loss_func(self):
     @abstractmethod
     def loss_func(self, loss_mask: torch.Tensor, output_tensor: torch.Tensor):
-        raise NotImplementedError
+        pass
 
-    # def get_forward_step_func(self):
     @abstractmethod
     def forward_step(self, data_iterator, model: GPTModel):
-        raise NotImplementedError
+        pass

primus/modules/trainer/megatron/sft_trainer.py

@@ -6,14 +6,11 @@ def __init__(self, *args, **kwargs):
         kwargs["module_name"] = "sft_trainer"
         super().__init__(*args, **kwargs)
 
-    def get_batch_func(self):
+    def get_batch(self, data_iterator):
         raise NotImplementedError
 
-    def get_loss_func(self):
+    def loss_func(self, loss_mask: torch.Tensor, output_tensor: torch.Tensor):
         raise NotImplementedError
 
-    def build_dataset_and_tokenizer(self):
-        raise NotImplementedError
-
-    def get_forward_step_func(self):
+    def forward_step(self, data_iterator, model: GPTModel):
         raise NotImplementedError

primus/modules/trainer/megatron/trainer.py

@@ -2,7 +2,6 @@
 import os
 import time
 from contextlib import nullcontext
-from datetime import datetime
 from typing import Union
 
 import megatron
@@ -1721,7 +1720,8 @@ def training_log(
                     writer.add_scalar("iteration-time", elapsed_time_per_iteration, iteration)
                 if wandb_writer:
                     wandb_writer.log({"iteration-time": elapsed_time_per_iteration}, iteration)
-            log_string = f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]"
+            # log_string = f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]"
+            log_string = f""
             if hasattr(self, "episode_count") and self.episode_count is not None:
                 log_string += f" episode {self.episode_count} |"
             log_string += " iteration {:8d}/{:8d} |".format(iteration, args.train_iters)