Skip to content

Commit a59a564

Browse files
balasaajaybalasaajay
authored
Merge branch 'main' into aut-498
2 parents ad2e4e1 + 4415119 commit a59a564

86 files changed

Lines changed: 120814 additions & 618 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎.github/actions/action.yml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,10 @@ inputs:
6565
description: "Trigger cadence for cadence filter (pr|nightly|mergegroup). Empty disables filter."
6666
required: false
6767
default: ""
68+
sha:
69+
description: "Git ref to check out. Must match the SHA used by the upstream parse step so recipes don't diverge between scheduling and execution."
70+
required: false
71+
default: ""
6872
runs:
6973
using: "composite"
7074
steps:
@@ -74,6 +78,8 @@ runs:
7478

7579
- name: Checkout repository
7680
uses: actions/checkout@v6
81+
with:
82+
ref: ${{ inputs.sha }}
7783

7884
- name: Change ownership of /home/runner/
7985
shell: bash

‎.github/workflows/cicd-main.yml‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -709,6 +709,7 @@ jobs:
709709
is_unit_test: "true"
710710
PAT: ${{ secrets.PAT }}
711711
container-image: ${{ env.container-registry }}/megatron-lm:${{ needs.configure.outputs.sha }}
712+
sha: ${{ needs.configure.outputs.sha }}
712713

713714
cicd-parse-integration-tests-h100:
714715
runs-on: ubuntu-latest
@@ -728,7 +729,7 @@ jobs:
728729
success()
729730
|| needs.pre-flight.outputs.is_ci_workload == 'true'
730731
|| needs.pre-flight.outputs.force_run_all == 'true'
731-
|| needs.pre-flight.outputs.is_merge_group == 'true'
732+
|| (needs.pre-flight.outputs.is_merge_group == 'true' && !failure())
732733
)
733734
&& !cancelled()
734735
outputs:
@@ -805,7 +806,7 @@ jobs:
805806
success()
806807
|| needs.pre-flight.outputs.is_ci_workload == 'true'
807808
|| needs.pre-flight.outputs.force_run_all == 'true'
808-
|| needs.pre-flight.outputs.is_merge_group == 'true'
809+
|| (needs.pre-flight.outputs.is_merge_group == 'true' && !failure())
809810
)
810811
&& !cancelled()
811812
steps:
@@ -827,6 +828,7 @@ jobs:
827828
n_repeat: ${{ needs.configure.outputs.n_repeat }}
828829
lightweight: ${{ needs.configure.outputs.lightweight }}
829830
cadence: ${{ needs.configure.outputs.cadence }}
831+
sha: ${{ needs.configure.outputs.sha }}
830832

831833
cicd-parse-integration-tests-gb200:
832834
runs-on: ubuntu-latest
@@ -849,7 +851,7 @@ jobs:
849851
success()
850852
|| needs.pre-flight.outputs.is_ci_workload == 'true'
851853
|| needs.pre-flight.outputs.force_run_all == 'true'
852-
|| needs.pre-flight.outputs.is_merge_group == 'true'
854+
|| (needs.pre-flight.outputs.is_merge_group == 'true' && !failure())
853855
)
854856
&& !cancelled()
855857
outputs:
@@ -928,7 +930,7 @@ jobs:
928930
success()
929931
|| needs.pre-flight.outputs.is_ci_workload == 'true'
930932
|| needs.pre-flight.outputs.force_run_all == 'true'
931-
|| needs.pre-flight.outputs.is_merge_group == 'true'
933+
|| (needs.pre-flight.outputs.is_merge_group == 'true' && !failure())
932934
)
933935
&& !cancelled()
934936
steps:
@@ -951,6 +953,7 @@ jobs:
951953
lightweight: ${{ needs.configure.outputs.lightweight }}
952954
platform: dgx_gb200
953955
cadence: ${{ needs.configure.outputs.cadence }}
956+
sha: ${{ needs.configure.outputs.sha }}
954957

955958
Nemo_CICD_Test:
956959
needs:

‎gpt_builders.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,7 @@ def _get_transformer_layer_spec(use_te, config):
123123
use_kitchen_attention=config.use_kitchen_attention,
124124
kitchen_attention_backend=config.kitchen_attention_backend,
125125
mla_down_proj_fusion=getattr(config, "mla_down_proj_fusion", False),
126+
use_grouped_gemm_for_dense_mlp=config.use_grouped_gemm_for_dense_mlp,
126127
)
127128
elif config.transformer_impl == "inference_optimized":
128129
return get_gpt_layer_with_inference_spec(

‎megatron/core/_rank_utils.py‎

Lines changed: 50 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,16 +4,22 @@
44

55
import logging
66
import os
7+
import warnings
78
from typing import Any
89

910
import torch
1011

12+
from megatron.core._slurm_utils import resolve_slurm_rank, resolve_slurm_world_size
13+
1114

1215
def safe_get_rank() -> int:
13-
"""Safely get the rank of the current process.
16+
"""Get the distributed rank safely, even if torch.distributed is not initialized.
1417
15-
Returns the rank from torch.distributed if initialized, otherwise falls back
16-
to the RANK environment variable, defaulting to 0.
18+
Fallback order:
19+
1. torch.distributed.get_rank() (if initialized)
20+
2. RANK environment variable (torchrun/torchelastic)
21+
3. SLURM_PROCID environment variable (SLURM)
22+
4. Default: 0 (with warning)
1723
1824
Returns:
1925
int: The rank of the current process.
@@ -23,11 +29,51 @@ def safe_get_rank() -> int:
2329

2430
# If torch.distributed is not initialized, try to read environment variables.
2531
try:
26-
return int(os.environ.get("RANK", 0))
32+
if "RANK" in os.environ:
33+
return int(os.environ["RANK"])
34+
35+
slurm_rank = resolve_slurm_rank()
36+
if slurm_rank is not None:
37+
return slurm_rank
38+
39+
warnings.warn(
40+
"Could not determine rank from torch.distributed, RANK, or SLURM_PROCID. "
41+
"Defaulting to rank 0."
42+
)
43+
return 0
2744
except (ValueError, TypeError):
2845
return 0
2946

3047

48+
def safe_get_world_size() -> int:
49+
"""Get the distributed world size safely, even if torch.distributed is not initialized.
50+
51+
Fallback order:
52+
1. torch.distributed.get_world_size() (if initialized)
53+
2. WORLD_SIZE environment variable (torchrun/torchelastic)
54+
3. SLURM_NTASKS environment variable (SLURM)
55+
4. Default: 1 (with warning)
56+
57+
Returns:
58+
The total number of processes in the distributed job.
59+
"""
60+
if torch.distributed.is_initialized():
61+
return torch.distributed.get_world_size()
62+
63+
if "WORLD_SIZE" in os.environ:
64+
return int(os.environ["WORLD_SIZE"])
65+
66+
slurm_world_size = resolve_slurm_world_size()
67+
if slurm_world_size is not None:
68+
return slurm_world_size
69+
70+
warnings.warn(
71+
"Could not determine world size from torch.distributed, WORLD_SIZE, or SLURM_NTASKS. "
72+
"Defaulting to world size 1."
73+
)
74+
return 1
75+
76+
3177
def log_single_rank(logger: logging.Logger, *args: Any, rank: int = 0, **kwargs: Any) -> None:
3278
"""Log a message only on a single rank.
3379

‎megatron/core/_slurm_utils.py‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
2+
3+
"""Utilities for detecting and configuring SLURM cluster environments.
4+
5+
This module provides functionality to detect SLURM environments and extract
6+
distributed training configuration from SLURM environment variables.
7+
"""
8+
9+
import os
10+
11+
12+
def is_slurm_job() -> bool:
13+
"""Detect if running in a SLURM environment.
14+
15+
Returns:
16+
True if SLURM job detected, False otherwise.
17+
"""
18+
return "SLURM_NTASKS" in os.environ
19+
20+
21+
def resolve_slurm_rank() -> int | None:
22+
"""Get the global rank from SLURM environment.
23+
24+
Returns:
25+
The global rank, or None if not in SLURM environment.
26+
"""
27+
if not is_slurm_job():
28+
return None
29+
return int(os.environ["SLURM_PROCID"]) if "SLURM_PROCID" in os.environ else None
30+
31+
32+
def resolve_slurm_world_size() -> int | None:
33+
"""Get the world size from SLURM environment.
34+
35+
Returns:
36+
The world size, or None if not in SLURM environment.
37+
"""
38+
if not is_slurm_job():
39+
return None
40+
return int(os.environ["SLURM_NTASKS"]) if "SLURM_NTASKS" in os.environ else None
41+
42+
43+
def resolve_slurm_local_rank() -> int | None:
44+
"""Get the local rank from SLURM environment.
45+
46+
Returns:
47+
The local rank, or None if not in SLURM environment.
48+
"""
49+
if not is_slurm_job():
50+
return None
51+
return int(os.environ["SLURM_LOCALID"]) if "SLURM_LOCALID" in os.environ else None

‎megatron/core/distributed/distributed_data_parallel.py‎

Lines changed: 3 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77
import torch
88

99
from ..config_logger import has_config_logger_enabled, log_config_to_disk
10-
from ..fp8_utils import is_float8tensor, post_all_gather_processing
1110
from ..optimizer.param_layout import FullParamLayout
1211
from ..process_groups_config import ProcessGroupCollection
1312
from ..transformer.cuda_graphs import is_graph_capturing
@@ -476,53 +475,20 @@ def no_sync(self):
476475
def _start_bucket_group_param_sync(
477476
self, bucket_group: '_ParamAndGradBucketGroup', force_sync: bool
478477
) -> None:
479-
"""Dispatch one bucket group's param all-gather + run the FP8 / MXFP8
478+
"""Dispatch one bucket group's param all-gather + run the FP8 / MXFP8 / FP4
480479
post-all-gather work the synchronous path needs.
481480
482481
Factored out of :meth:`start_param_sync` so callers that own a subset
483482
of bucket groups (e.g. a chained ``LayerWiseDistributedOptimizer`` +
484483
``DistributedOptimizer`` pair) can sync only their own buckets without
485-
losing the FP8 post-processing that follows the collective.
484+
losing the post-processing that follows the collective.
486485
"""
487486
bucket_group.start_param_sync(force_sync=force_sync)
488487

489488
if self.ddp_config.overlap_param_gather:
490489
return
491490

492-
# For MXFP8 params, we need to copy the all-gathered param data from the buffer to
493-
# the param.data, since param buffer is not mapped to model params for MXFP8 case.
494-
# The paramaters are cast from bf16 to MXFP8 during copy.
495-
# In the case of "overlap_param_gather=True", the param copy is done
496-
# in "finish_param_sync" stage after zeroing the shared gardient buffers.
497-
if self.ddp_config.reuse_grad_buf_for_mxfp8_param_ag:
498-
for bucket in bucket_group.buckets:
499-
is_bf16_weight_bucket = False
500-
for param in bucket.params:
501-
# Skip copying since bf16 weights in the mxfp8 model
502-
# are already mapped to param.data.
503-
if not is_float8tensor(param):
504-
is_bf16_weight_bucket = True
505-
break
506-
param_start, param_end = bucket.param_to_index[param]
507-
param_slice = bucket.param_data.view(-1)[param_start:param_end]
508-
param.data.copy_(param_slice.view(param.data.shape))
509-
if is_bf16_weight_bucket:
510-
continue
511-
# All-gathered params are not needed after being copied to param.data.
512-
# Zero out the param buffer (shared with grad buffer) for gradient
513-
# accumulation. We cannot zero out the entire grad buffer because one grad
514-
# buffer may correspond to multiple param buffers. If we zero out the entire
515-
# grad buffer, it would clear the data of those param buffers that have not
516-
# yet completed AG.
517-
bucket.param_data.zero_()
518-
else:
519-
fp8_params = []
520-
for bucket in bucket_group.buckets:
521-
for param in bucket.params:
522-
if is_float8tensor(param):
523-
fp8_params.append(param)
524-
if len(fp8_params) > 0:
525-
post_all_gather_processing(fp8_params)
491+
bucket_group._post_param_sync()
526492

527493
def start_param_sync(self, *unused, force_sync: bool = False, force_dispatch: bool = False):
528494
"""

‎megatron/core/distributed/param_and_grad_buffer.py‎

Lines changed: 39 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -267,6 +267,38 @@ def reset(self):
267267
self.per_param_grad_ready_counts = {}
268268
self.is_last_microbatch = True
269269

270+
def _post_param_sync(self):
271+
"""Run post-processing after param all-gather completes."""
272+
if self.ddp_config.reuse_grad_buf_for_mxfp8_param_ag:
273+
for bucket in self.buckets:
274+
is_bf16_weight_bucket = False
275+
for param in bucket.params:
276+
# Skip copying since bf16 weights in the mxfp8 model
277+
# are already mapped to param.data.
278+
if not is_float8tensor(param):
279+
is_bf16_weight_bucket = True
280+
break
281+
param_start, param_end = bucket.param_to_index[param]
282+
param_slice = bucket.param_data.view(-1)[param_start:param_end]
283+
param.data.copy_(param_slice.view(param.data.shape))
284+
if is_bf16_weight_bucket:
285+
continue
286+
# All-gathered params are not needed after being copied to param.data.
287+
# Zero out the param buffer (shared with grad buffer) for gradient accumulation.
288+
# We cannot zero out the entire grad buffer because one grad buffer may
289+
# correspond to multiple param buffers. If we zero out the entire grad buffer,
290+
# it would clear the data of those param buffers that have not yet completed AG.
291+
bucket.param_data.zero_()
292+
return
293+
294+
quantized_params = []
295+
for bucket in self.buckets:
296+
for param in bucket.params:
297+
if is_float8tensor(param) or is_nvfp4tensor(param):
298+
quantized_params.append(param)
299+
if len(quantized_params) > 0:
300+
post_all_gather_processing(quantized_params)
301+
270302
def check_grads(self, check_for_nan_or_inf, check_for_large):
271303
"""
272304
Make sure norm of grads in bucket are not NaN prior to data-parallel
@@ -325,6 +357,7 @@ def start_param_sync(self, force_sync: bool = False):
325357
if self.param_gather_handle is not None:
326358
self.param_gather_handle.wait()
327359
self.param_gather_handle = None
360+
self._post_param_sync()
328361
return
329362
else:
330363
assert self.param_gather_handle is None
@@ -344,6 +377,8 @@ def start_param_sync(self, force_sync: bool = False):
344377
dp_size = self.intra_distributed_optimizer_instance_size
345378
if dp_size == 1:
346379
# Single-rank group (e.g., expt_dp_size == 1): no all-gather needed.
380+
if force_sync and self.ddp_config.overlap_param_gather:
381+
self._post_param_sync()
347382
self.param_gather_dispatched = True
348383
return
349384
local_rank = self.intra_distributed_optimizer_instance_rank
@@ -441,6 +476,8 @@ def start_param_sync(self, force_sync: bool = False):
441476
# (async_op=False) is used, `cm` is not None. Manually set to None for
442477
# consistency with prior code.
443478
self.param_gather_handle = None
479+
if force_sync and self.ddp_config.overlap_param_gather:
480+
self._post_param_sync()
444481
self.param_gather_dispatched = True
445482

446483
def finish_param_sync(self, skip_next_bucket_dispatch: bool = False):
@@ -480,30 +517,7 @@ def finish_param_sync(self, skip_next_bucket_dispatch: bool = False):
480517
else:
481518
self.next_param_gather_bucket_group.start_param_sync()
482519

483-
# For the mxfp8_param with "reuse_grad_buf_for_mxfp8_param_ag=True",
484-
# we need to copy the param_data from the shared_param/grad_buffer to param.data
485-
# after the param all-gather.
486-
if self.ddp_config.reuse_grad_buf_for_mxfp8_param_ag:
487-
for bucket in self.buckets:
488-
is_bf16_weight_bucket = False
489-
for param in bucket.params:
490-
# Skip copying since bf16 weights in the mxfp8 model
491-
# are already mapped to param.data.
492-
if not is_float8tensor(param):
493-
is_bf16_weight_bucket = True
494-
break
495-
param_start, param_end = bucket.param_to_index[param]
496-
param_slice = bucket.param_data.view(-1)[param_start:param_end]
497-
param.data.copy_(param_slice.view(param.data.shape))
498-
if is_bf16_weight_bucket:
499-
continue
500-
# All-gathered params are not needed after being copied to param.data.
501-
# Zero out the param buffer (shared with grad buffer) for gradient accumulation.
502-
# We cannot zero out the entire grad buffer because one grad buffer may
503-
# correspond to multiple param buffers. If we zero out the entire grad buffer,
504-
# it would clear the data of those param buffers that have not yet completed AG.
505-
bucket.param_data.zero_()
506-
elif not self.ddp_config.use_distributed_optimizer:
520+
if not self.ddp_config.use_distributed_optimizer:
507521
for bucket in self.buckets:
508522
if bucket.layerwise_gather_list is None:
509523
continue
@@ -526,14 +540,7 @@ def finish_param_sync(self, skip_next_bucket_dispatch: bool = False):
526540
# (a view into grad_data) would start from the result of the
527541
# latest parameter all-gather instead of zero.
528542
bucket.grad_data.zero_()
529-
else:
530-
fp8_params = []
531-
for bucket in self.buckets:
532-
for param in bucket.params:
533-
if is_float8tensor(param):
534-
fp8_params.append(param)
535-
if len(fp8_params) > 0:
536-
post_all_gather_processing(fp8_params)
543+
self._post_param_sync()
537544

538545
def start_grad_sync(self, force_all_reduce: Optional[bool] = False):
539546
"""

0 commit comments

Comments
 (0)