Upgrade: Sync to Megatron-LM Core 0.17.0#32
Closed
lxd-cumt wants to merge 932 commits into
Closed
Conversation
Signed-off-by: oliver könig <okoenig@nvidia.com>
Signed-off-by: oliver könig <okoenig@nvidia.com>
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
Signed-off-by: oliver könig <okoenig@nvidia.com>
Signed-off-by: oliver könig <okoenig@nvidia.com>
…`. (NVIDIA#2639) Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>
Signed-off-by: oliver könig <okoenig@nvidia.com>
Signed-off-by: Lifu Zhang <lifuz@login-lyris01.lyris.clusters.nvidia.com> Signed-off-by: Lifu Zhang <lifuz@login-lyris02.lyris.clusters.nvidia.com> Co-authored-by: Lifu Zhang <lifuz@login-lyris01.lyris.clusters.nvidia.com> Co-authored-by: Lifu Zhang <lifuz@login-lyris02.lyris.clusters.nvidia.com>
Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>
Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>
Co-authored-by: Eric Harper <complex451@gmail.com> Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>
Signed-off-by: dimapihtar <dpihtar@gmail.com>
Co-authored-by: Eric Harper <eharper@nvidia.com>
…ors in core 0.17.0 - fp8_utils.py: import cur_platform from megatron.plugin.platform instead of megatron.core.utils - optimizer/__init__.py: remove stale MuonOptimizerConfig import (class doesn't exist)
…p engram parallel - Restore muon.py and optimizer_config.py from upstream core_r0.17.0 - Restore optimizer/__init__.py from upstream, re-apply engram parallel patches - Upstream now has full Muon support natively
…_v0.17.0 functions - Add global declaration for _LAST_RANK_WHEN_USING_PIPELINE in parallel_state.py - Add is_dualpipev_first_chunk parameter to get_num_layers_to_build() and get_transformer_layer_offset() - Add is_mtp_layer parameter to TopKRouter.__init__() - Fix mtp_model_layer_spec keyword in get_mtp_layer_spec_for_backend() call These parameters are FlagScale-specific extensions layered on top of upstream Megatron-LM core_v0.17.0 to support dualpipev and MTP features. Verified with DeepSeek V3 16B training (PP=2, MoE, MTP enabled).
…oups The arg was renamed from --disable-gloo-process-groups (dest=enable_gloo_process_groups) to use_gloo_process_groups in common_config.py. Update the reference in hetero ProcessMesh to match. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Function lives in megatron.training.training, not megatron.training.utils. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…m merge Engram embedding params (is_engram_embedding=True, allreduce=False) were landing in expert_parallel_buffers after the upstream merge dropped the FlagScale DDP patch. This caused KeyError in DistributedOptimizer because the MoE param group filter excluded engram params but the MoE buffers still contained them. Restored patches: - engram_dp_group from process_group_dict - engram_embedding_params split from expert_parallel_params - engram_embedding_buffers/bucket_groups allocation - engram buffers included in all iteration loops - broadcast_params routes engram params to engram_dp_group
Restored patches: - distributed_data_parallel.py: engram DDP buffer separation - optimizer/__init__.py: dense/MoE/engram param group filters, mp_group list wrapping - optimizer_config.py: vision_ration field for grouped learning rate - parallel_state.py: engram/dualpipev/hetero parallel state functions - p2p_communication.py: hetero PP group parameter - schedules.py: dualpipev dispatch in get_forward_backward_func - utils.py: engram embedding data utilities
- transformer_block.py: refined recompute save/restore pattern - utils.py: parse_te_version_str() for TE-FL version strings - schedules.py: cur_platform device abstraction replacing hardcoded cuda - model_chunk_schedule_plan.py: engram fine-grained scheduling support - fine_grained_callables.py: engram forward hook in submodule_attn_forward - clip_grads.py: @overridable decorators + cur_platform import
Restore FlagScale plugin platform abstraction that was lost during upstream 0.17.0 sync. Replaces hardcoded torch.cuda.* calls with cur_platform.*() for multi-platform hardware support (CUDA/MUSA/TXDA/CPU). Files modified: - distributed: distributed_data_parallel.py, param_and_grad_buffer.py, fsdp/param_and_grad_buffer.py - optimizer: distrib_optimizer.py (also migrated type() checks to device.type + dtype) - parallel_state.py, utils.py, rerun_state_machine.py - pipeline_parallel: bridge_communicator.py, utils.py - transformer: attention.py, transformer_layer.py, multi_token_prediction.py, moe_utils.py - models: gpt_model.py, model_chunk_schedule_plan.py, fine_grained_callables.py - full_cuda_graph.py (kept torch.cuda.CUDAGraph/graph as CUDA-specific) - inference: dynamic_context.py, dynamic_engine.py, mamba_metadata.py, abstract_model_inference_wrapper.py - dist_checkpointing: strategies/torch.py
Upstream changes (10 commits since 9539a12): - NVFP4 native weights support for DDP (NVIDIA#4290) - Bump mfsdp to 0.4.0 - Documentation and version updates Overlapping files (upstream + cur_platform patches preserved): - distributed_data_parallel.py: added NVFP4 support - param_and_grad_buffer.py: NVFP4 dual-buffer layout - distrib_optimizer.py: FP4 param handling New upstream files taken as-is: - fp4_utils.py, test_fp4_param.py, distributed_data_parallel_config.py Excluded: megatron/training/arguments.py (managed by FlagScale)
…latform patches - Replace old dev branch code with upstream fa49087 for: - megatron/core/extensions/kitchen.py - megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py - megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py - megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py - megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py - megatron/core/distributed/fsdp/mcore_fsdp_adapter.py - Apply cur_platform patches (torch.cuda -> cur_platform) for multi-platform support - kitchen.py and uneven_dtensor.py match upstream exactly (no torch.cuda calls)
Restore upstream code for files that had gratuitous changes: - language_module.py: restore MTP logic removed during sync - router.py: restore MTP logic removed during sync - cuda_graphs.py: remove stray wgrads_with_placeholders block - optimizer.py: remove junk 'from re import M' import, revert weights_only=False - schedule_plan.py: restore upstream formatting for TransformerLayerSchedulePlan call - QuickStart.md: restore upstream formatting - mamba_builders.py: restore removed vp_stage param - pretrain_mamba.py: restore removed loss_mask param All FlagScale-specific patches (cur_platform, @overridable, engram, is_offloading_candidate) preserved.
- setup_embeddings_and_output_layer: add hasattr guards for embedding/output_layer, add MuP is_embedding_parameter marking, add init_model_with_meta_device guard - get_grad_norm_fp32: import multi_tensor_scale_tensor_impl, add conditional tensor vs scalar path for total_norm computation - count_zeros_fp32: fix dtype from torch.float to torch.int64
FP4 param gather modifications should be on FlagScale side, not Megatron-LM-FL.
- Remove accidental 'from re import M' in shared_experts.py - Revert is_offloading_candidate check in hybrid_optimizer.py to upstream logic (attribute never set on any param)
These files are local-only and should not be tracked in git.
…17.0 These directories do not need cur_platform patches and should match upstream exactly.
Used by FlagScale engram multi_head_embedding to mark weights for CPU offloading.
Move upstream imports outside FlagScale Begin/End blocks in: - clip_grads.py - distrib_optimizer.py - distributed_data_parallel.py - parallel_state.py - schedules.py
When qk_layernorm_hidden_dim is True, use query_projection_size/kv_projection_size instead of hidden_size_per_attention_head for QK layernorm.
When megatron-core is pip-installed (non-editable), the Makefile is not included in the wheel but the pre-compiled .so is. compile_helpers() now checks for the .so before attempting make, avoiding a crash in CI/CD.
Custom TE builds (e.g. v0.1.0+te2.9.0) may report >= 2.5.0 via the +te version tag but lack the retain_pinned_cpu_buffers parameter in get_cpu_offload_context. Use inspect.signature to check the actual function before passing the kwarg.
…rallelConfig The plugin finalize_model_grads.py references this field but it was missing from the upstream config dataclass after the sync.
|
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
20 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.