[megatron] Support deepseek-v4 megatron

[Jintao-Huang]

modelscope/mcore-bridge#86

[gemini-code-assist]

Code Review

This pull request updates the documentation to indicate support for DeepSeek-V4 models and adds new configuration arguments to the MegatronArguments class, such as csa_dense_mode and use_fused_mhc. A suggestion was made to improve code readability by adding a blank line before the new DeepSeek-V4 argument section in megatron_args.py to maintain consistent formatting.

[gemini-code-assist]

Add a blank line before the # deepseek-v4 comment to maintain consistent section separation and improve readability. The previous section (dsa) ends at line 626, and a blank line would help distinguish the new group of arguments, following the style used for the # other section below.

Suggested change

	# deepseek-v4
	# deepseek-v4

[Jintao-Huang]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces support for DeepSeek-V4 models, including documentation updates and the addition of specific Megatron arguments such as csa_dense_mode and use_fused_mhc. Changes were also made to pipeline parallel batch processing and model conversion logic. Feedback highlights that removing the memory optimization for input_ids on intermediate pipeline stages should be conditional to avoid overhead for other models. Furthermore, the use of hardcoded .cuda() calls in model conversion should be replaced with device-agnostic logic to ensure portability across different hardware.

I am having trouble creating individual review comments. Click here to see my feedback.

swift/megatron/trainers/utils.py (28-29)

The removal of the logic that nulls input_ids on non-first pipeline stages (when MTP is not used) will increase memory consumption across all pipeline ranks for all models, as input_ids will now be retained on every rank. If DeepSeek-V4 requires input_ids on intermediate stages (e.g., for MHC or MTP), this should be handled conditionally (for example, by checking for DeepSeek-V4 specific flags like use_fused_mhc) to preserve the memory optimization for other models.

swift/megatron/utils/convert_utils.py (255-257)

Hardcoding .cuda() limits the portability of the code to other hardware accelerators (e.g., XPU, MPS). It is recommended to use .to(device) where the device is determined from the model. Additionally, using getattr to safely access mg_language_model.decoder is safer as this attribute may not exist on all model architectures supported by this utility.

    if getattr(config, 'enable_hyper_connections', False):
        decoder = getattr(mg_language_model, 'decoder', None)
        if decoder is not None:
            device = next(mg_language_model.parameters()).device
            for param in decoder.parameters(recurse=False):
                param.data = param.data.to(device)

[Jintao-Huang]

/gemini review

[gemini-code-assist]

Code Review

This pull request adds support for DeepSeek-V4 models, updates documentation, bumps dependency requirements (including upgrading the minimum mcore-bridge version to 1.3.0 and expanding the transformers version range), and introduces new Megatron arguments for DeepSeek-V4. Review feedback suggests aligning the transformers version constraint in requirements/install_all.sh with requirements/framework.txt and guarding a .cuda() call in convert_utils.py with torch.cuda.is_available() to prevent failures in CPU-only environments.