Skip to content

llama31 405b gb200 nvfp4 no pg overlap #3713

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
malay-nagda merged 2 commits into from
May 11, 2026
+14 −0
Merged

llama31 405b gb200 nvfp4 no pg overlap #3713

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
151d3dc
Disable param gather overlap for Llama3.1 405B GB200 NVFP4
malay-nagda May 6, 2026
a8a6ed4
doc string
malay-nagda May 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions scripts/performance/configs/llama/llama31_llm_pretrain.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,18 @@ def set_llama31_common_configs(cfg: ConfigContainer) -> None:
cfg.ddp.grad_reduce_in_fp32 = False


def disable_param_gather_overlap(cfg: ConfigContainer) -> None:
"""
Disable parameter-gather overlap to reduce training peak memory and avoid OOM.
Note: This is a workaround and should be removed once the issue is fixed.
See: https://github.com/NVIDIA-NeMo/Megatron-Bridge/issues/3714
"""
cfg.ddp.overlap_param_gather = False
cfg.optimizer.overlap_param_gather = False
cfg.comm_overlap.overlap_param_gather = False
cfg.comm_overlap.align_param_gather = False


def llama31_405b_pretrain_config_gb300(
precision: str = "bf16", mock: bool = True, config_variant: str = "v1"
) -> ConfigContainer:
Expand Down Expand Up @@ -108,6 +120,8 @@ def llama31_405b_pretrain_config_gb200(

cfg.comm_overlap.tp_comm_overlap_cfg = comm_overlap_cfg
cfg.comm_overlap.tp_comm_overlap = False if precision == "nvfp4" else cfg.comm_overlap.tp_comm_overlap
if precision == "nvfp4" and config_variant.lower() == "v2":
disable_param_gather_overlap(cfg)
Copy link
Copy Markdown
Contributor

@yaoyu-33 yaoyu-33 May 7, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@malay-nagda can you comment why need disable?

Copy link
Copy Markdown
Contributor Author

@malay-nagda malay-nagda May 11, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@yaoyu-33 added to doc string- https://github.com/NVIDIA-NeMo/Megatron-Bridge/pull/3713/changes#diff-549b27caa8a6b59ed45c30e07ff12a5981dd32f1d2b7b84f7bbca7186c9c9d7fR48


return cfg

Expand Down
Loading