update

Author: Jintao-Huang

swift/megatron/convert.py

@@ -74,7 +74,7 @@ def convert_hf2mcore(args: ExportArguments) -> None:
         logger.info(f'Successfully saved Megatron model weights in `{args.output_dir}`.')
     # Place it at the end to avoid test_convert_precision affecting precision.
     if args.test_convert_precision:
-        test_convert_precision(hf_model, mg_model, template, args.test_convert_dtype)
+        test_convert_precision(megatron_args, hf_model, mg_model, template, test_convert_dtype=args.test_convert_dtype)
 
 
 def convert_mcore2hf(args: ExportArguments) -> None:
@@ -131,7 +131,7 @@ def convert_mcore2hf(args: ExportArguments) -> None:
         logger.info(f'Successfully saved HF model weights in `{args.output_dir}`.')
         if args.test_convert_precision:
             hf_model, template = prepare_model_template(args, model=args.output_dir)
-            test_convert_precision(hf_model, mg_model, template, args.test_convert_dtype)
+            test_convert_precision(megatron_args, hf_model, mg_model, template, test_convert_dtype=args.test_convert_dtype)
     elif args.to_mcore:
         if args.thread_count is None:
             checkpoint_size = sum(get_n_params_grads(mg_model)[0]) * torch.finfo(args.torch_dtype).bits // 8e9

swift/megatron/init.py

@@ -518,6 +518,7 @@ def sharded_state_dict(
 def _patch_TransformerLayer():
     import megatron.core
     from megatron.core.transformer import TransformerLayer
+    from megatron.core import mpu
     _origin_forward = TransformerLayer.forward
     mcore_013 = version.parse(megatron.core.__version__) >= version.parse('0.13.0rc0')
 
@@ -531,7 +532,7 @@ def forward(self, *_args, **kwargs):
         if not mcore_013:
             return _origin_forward(self, *_args, **kwargs)
         hidden_states, context = self._forward_attention(*_args, **kwargs)
-        args = get_args()
+        args = self.config.args
         mlp_padding_free = args.mlp_padding_free and 'attention_mask' in kwargs
         mask = None
         if mlp_padding_free and hidden_states.shape[1] > 1:
@@ -660,6 +661,7 @@ def _write_item(self, *args, **kwargs):
 def _patch_mrope():
     from megatron.core.models.common.embeddings.rotary_pos_embedding import MultimodalRotaryEmbedding
     import megatron.core
+    from megatron.core import mpu
     from megatron.core.models.common.embeddings.rope_utils import _apply_rotary_pos_emb_bshd
     from megatron.core.models.common.embeddings import rope_utils
 
@@ -696,7 +698,7 @@ def forward(self, position_ids, mrope_section: List[int], packed_seq: bool = Fal
         seq_expanded = seq[:, :, None, :].float()
         # shape (3, bs, seq_length, dim)
         freqs = (inv_freq_expanded @ seq_expanded).transpose(2, 3)
-        args = get_args()
+        args = self.config.args
         if args.mrope_interleaved:
             freqs = apply_interleaved_mrope(freqs, mrope_section)
             emb = torch.cat((freqs, freqs), dim=-1)
@@ -744,8 +746,7 @@ def _apply_rotary_pos_emb_thd(
         if cp_group is not None:
             cp_size = cp_group.size()
         else:
-            args = get_args()
-            cp_size = args.context_parallel_size
+            cp_size = mpu.get_context_parallel_world_size()
         cu_seqlens_for_batched = cu_seqlens // cp_size
         use_batched_rope = (freqs.dim() >= 1 and freqs.shape[0] == cu_seqlens_for_batched[-1]).item()
         if not use_batched_rope:

swift/megatron/pipelines/export/export.py

@@ -82,7 +82,7 @@ def convert_mcore2hf(self) -> None:
                 device_map = args.device_map or 'auto'
                 hf_model, template = prepare_model_template(
                     args, device_map=device_map, **kwargs) if is_last_rank() else (None, template)
-            test_convert_precision(hf_model, mg_model, template, args.test_convert_dtype)
+            test_convert_precision(args, hf_model, mg_model, template, test_convert_dtype=args.test_convert_dtype)
             dist.barrier()
 
     def convert_hf2mcore(self) -> None:
@@ -135,7 +135,7 @@ def convert_hf2mcore(self) -> None:
                     device_map = args.device_map or 'auto'
                     hf_model, template = prepare_model_template(
                         args, device_map=device_map) if is_last_rank() else (None, template)
-                test_convert_precision(hf_model, mg_model, template, args.test_convert_dtype)
+                test_convert_precision(args, hf_model, mg_model, template, test_convert_dtype=args.test_convert_dtype)
                 dist.barrier()
             else:
                 logger.warning('Skip test_convert_precision because `--adapter_load` is specified.')

swift/megatron/trainers/gkd_trainer.py

@@ -441,7 +441,7 @@ def _compute_teacher_logits(self, encoded_batches: List[Dict], vp_stage: Optiona
             teacher_data.pop('labels', None)
             # Teacher forward with args override for correct hidden_size
             with self.load_teacher_model_context(), self._teacher_args_context(), torch.no_grad():
-                teacher_logits = forward_step_helper(teacher_model, teacher_data)
+                teacher_logits = forward_step_helper(self.args, teacher_model, teacher_data)
                 if teacher_logits is not None:
                     teacher_logits = teacher_logits.detach()
             encoded_batch['teacher_logits'] = teacher_logits

swift/megatron/trainers/grpo_trainer.py

@@ -1504,7 +1504,7 @@ def model_forward(self, model, data_iterator, no_grad=True, per_token=False):
         context = torch.no_grad() if no_grad else nullcontext()
 
         with context:
-            output_tensor = forward_step_helper(model, data)
+            output_tensor = forward_step_helper(self.args, model, data)
 
         # packed_seq_params only exists in padding_free mode
         packed_seq_params = data.get('packed_seq_params')

swift/megatron/utils/convert_utils.py

@@ -143,8 +143,9 @@ def get_examples(is_multimodal: bool) -> Dict[str, Any]:
     return data
 
 
-def test_convert_precision(args, hf_model, mg_model, template):
-    torch_dtype = args.test_convert_dtype
+def test_convert_precision(args, hf_model, mg_model, template, test_convert_dtype=None):
+    if test_convert_dtype is None:
+        test_convert_dtype = getattr(args, 'test_convert_dtype', torch.float32)
     template.set_mode('train')
     _test_params_sum(mg_model)
 
@@ -166,7 +167,7 @@ def test_convert_precision(args, hf_model, mg_model, template):
         ignore_modules = (model_arch.vision_tower + model_arch.aligner) if is_multimodal else []
         hf_modules = _find_modules(hf_model, ignore_modules=ignore_modules)
         with torch.inference_mode(), _model_cpu_forward_context(
-                hf_modules, torch_dtype, share_embedding=share_embedding):
+                hf_modules, test_convert_dtype, share_embedding=share_embedding):
             hf_inputs.pop('text_position_ids', None)
             hf_logits = hf_model(**hf_inputs).logits
             hf_logits = hf_logits.to('cuda')
@@ -195,8 +196,8 @@ def test_convert_precision(args, hf_model, mg_model, template):
             if n.endswith('router'):
                 m.to(mg_dtype)
     with torch.inference_mode(), _model_cpu_forward_context(
-            mg_modules, torch_dtype, 'cuda', share_embedding=share_embedding, target_device=mg_device):
-        mg_logits = forward_step_helper(mg_model, mg_inputs, dtype=torch_dtype)
+            mg_modules, test_convert_dtype, 'cuda', share_embedding=share_embedding, target_device=mg_device):
+        mg_logits = forward_step_helper(args, mg_model, mg_inputs, dtype=test_convert_dtype)
         if args.tensor_model_parallel_size > 1 and args.task_type != 'seq_cls':
             from megatron.core.tensor_parallel.mappings import gather_from_tensor_model_parallel_region
             if mg_logits is not None: