Skip to content

Clean up pretrain_gpt.py and pretrain_hybrid.py formatting and remove module globals #5351

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
ilml merged 3 commits into from
Jun 16, 2026
Merged

Clean up pretrain_gpt.py and pretrain_hybrid.py formatting and remove module globals #5351

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
bee7ca8
Clean up pretrain_gpt.py formatting and remove module globals
ilml Jun 16, 2026
184b3f9
Merge branch 'main' into cleanup/pretrain-gpt-globals
ilml Jun 16, 2026
0dc4317
Apply the same cleanup to pretrain_hybrid.py
ilml Jun 16, 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
157 changes: 120 additions & 37 deletions pretrain_gpt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", category=FutureWarning)

from functools import partial
from functools import lru_cache, partial
from typing import Any, List, Optional, Tuple

import torch
Expand Down Expand Up @@ -54,7 +54,7 @@
print_rank_0,
set_startup_timestamps,
)
from megatron.training.argument_utils import pretrain_cfg_container_from_args, gpt_config_from_args
from megatron.training.argument_utils import gpt_config_from_args, pretrain_cfg_container_from_args
from megatron.training.arguments import core_transformer_config_from_args, parse_and_validate_args
from megatron.training.datasets.fim_dataset import GPTFIMDataset, GPTFIMDatasetConfig
from megatron.training.datasets.sft_dataset import SFTDataset
Expand All @@ -72,20 +72,38 @@

stimer = StragglerDetector()

# Canonical, ordered schema of the fields ``get_batch`` returns. Kept alphabetical
# to match the historical ``sorted(batch.keys())`` order that callers unpack into.
BATCH_KEYS = [
"attention_mask",
"cu_seqlens",
"cu_seqlens_padded",
"hybrid_cp_group",
"labels",
"local_cp_size",
"loss_mask",
"max_seqlen",
"position_ids",
"tokens",
]


def get_batch(data_iterator, vp_stage: Optional[int] = None):
"""Generate a batch."""

BATCH_KEYS = ["attention_mask", "cu_seqlens", "cu_seqlens_padded", "hybrid_cp_group", "labels", "local_cp_size", "loss_mask", "max_seqlen", "position_ids", "tokens"]

args = get_args()
config = core_transformer_config_from_args(args)

cp_size = args.context_parallel_size
tp_rank = mpu.get_tensor_model_parallel_rank()
is_sft = args.sft
create_attention_mask_in_dataloader = args.create_attention_mask_in_dataloader
mtp_on_this_rank = mtp_on_this_rank_func(layout=config.pipeline_model_parallel_layout, mtp_num_layers=config.mtp_num_layers, ignore_virtual=False, vp_stage=vp_stage)
mtp_on_this_rank = mtp_on_this_rank_func(
layout=config.pipeline_model_parallel_layout,
mtp_num_layers=config.mtp_num_layers,
ignore_virtual=False,
vp_stage=vp_stage,
)
is_hybrid_cp = args.hybrid_context_parallel

if not is_first_or_last_pipeline_stage(vp_stage) and not mtp_on_this_rank and not is_sft:
Expand All @@ -95,15 +113,50 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None):
if tp_rank == 0:
batch = next(data_iterator)
for key in BATCH_KEYS:
batch[key] = batch[key].cuda(non_blocking=True) if key in batch and batch[key] is not None else None
batch[key] = (
batch[key].cuda(non_blocking=True)
if key in batch and batch[key] is not None
else None
)

batch = get_batch_on_this_tp_rank(
batch,
broadcast_src_rank=mpu.get_tensor_model_parallel_src_rank(),
broadcast_group=mpu.get_tensor_model_parallel_group(),
is_sft=is_sft,
is_hybrid_cp=is_hybrid_cp,
create_attention_mask_in_dataloader=create_attention_mask_in_dataloader,
cp_size=cp_size,
tp_rank=tp_rank,
micro_batch_size=args.micro_batch_size,
seq_length=args.seq_length,
mtp_on_this_rank=mtp_on_this_rank,
pipeline_model_parallel_size=args.pipeline_model_parallel_size,
is_pipeline_first_stage=mpu.is_pipeline_first_stage(),
is_pipeline_last_stage=mpu.is_pipeline_last_stage(),
)

batch = get_batch_on_this_tp_rank(batch, broadcast_src_rank=mpu.get_tensor_model_parallel_src_rank(), broadcast_group=mpu.get_tensor_model_parallel_group(), is_sft=is_sft, is_hybrid_cp=is_hybrid_cp, create_attention_mask_in_dataloader=create_attention_mask_in_dataloader, cp_size=cp_size, tp_rank=tp_rank, micro_batch_size=args.micro_batch_size, seq_length=args.seq_length, mtp_on_this_rank=mtp_on_this_rank, pipeline_model_parallel_size=args.pipeline_model_parallel_size, is_pipeline_first_stage=mpu.is_pipeline_first_stage(), is_pipeline_last_stage=mpu.is_pipeline_last_stage())

if not is_first_or_last_pipeline_stage(vp_stage) and not mtp_on_this_rank:
assert is_sft
return None, batch['cu_seqlens'], batch['cu_seqlens_padded'], None, None, None, None, batch['max_seqlen'], None, None

batch = get_batch_on_this_cp_rank(batch, is_hybrid_cp=is_hybrid_cp, cp_group=get_context_parallel_group(), hybrid_cp_group_func=get_hybrid_data_context_parallel_groups)
return (
None,
batch['cu_seqlens'],
batch['cu_seqlens_padded'],
None,
None,
None,
None,
batch['max_seqlen'],
None,
None,
)

batch = get_batch_on_this_cp_rank(
batch,
is_hybrid_cp=is_hybrid_cp,
cp_group=get_context_parallel_group(),
hybrid_cp_group_func=get_hybrid_data_context_parallel_groups,
)

# Return values in BATCH_KEYS order so callers can unpack into the fixed
# names regardless of any provenance fields wrappers like BlendedDataset
Expand All @@ -117,7 +170,27 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None):
# define spiky loss as a loss that's 10x the max loss observed
SPIKY_LOSS_FACTOR = 10

loss_func_cached_logits = None

@lru_cache(maxsize=1)
def _build_cached_logits_loss_func(
logprobs_dir, decode_threads, prefetch_factor, msc_prefetch_depth, kd_loss_alpha, ignore_errors
):
"""Build (once) the offline knowledge-distillation loss callable for cached logits.

Memoized so the teacher log-probability reader is constructed a single time per
process, replacing the previous module-level mutable global.
"""
from megatron.training.distillation import LossFuncCallable

return LossFuncCallable(
logprobs_dir=logprobs_dir,
decode_threads=decode_threads,
prefetch_factor=prefetch_factor,
msc_prefetch_depth=msc_prefetch_depth,
kd_loss_alpha=kd_loss_alpha,
ignore_errors=ignore_errors,
)


def loss_func(
loss_mask: torch.Tensor, output_tensor: torch.Tensor, model: Optional[GPTModel] = None
Expand All @@ -139,17 +212,14 @@ def loss_func(

if args.logits_load_dir is not None:
# Offline knowledge distillation loss using cached teacher log-probabilities.
global loss_func_cached_logits
if loss_func_cached_logits is None:
from megatron.training.distillation import LossFuncCallable
loss_func_cached_logits = LossFuncCallable(
logprobs_dir=args.logits_load_dir,
decode_threads=args.logits_load_decode_threads,
prefetch_factor=args.logits_load_prefetch_factor,
msc_prefetch_depth=args.logits_load_msc_prefetch_depth,
kd_loss_alpha=args.logits_load_kd_loss_alpha,
ignore_errors=args.logits_load_ignore_errors,
)
loss_func_cached_logits = _build_cached_logits_loss_func(
logprobs_dir=args.logits_load_dir,
decode_threads=args.logits_load_decode_threads,
prefetch_factor=args.logits_load_prefetch_factor,
msc_prefetch_depth=args.logits_load_msc_prefetch_depth,
kd_loss_alpha=args.logits_load_kd_loss_alpha,
ignore_errors=args.logits_load_ignore_errors,
)
loss, num_tokens, report = loss_func_cached_logits(loss_mask, output_tensor, model=model)
elif has_nvidia_modelopt and getattr(args, 'modelopt_enabled', False): # [ModelOpt]
loss, num_tokens, report = loss_func_modelopt(loss_mask, output_tensor, model=model)
Expand All @@ -168,14 +238,14 @@ def loss_func(
result=loss,
rejection_func=torch.isnan,
message="found NaN in local forward loss calculation",
tolerance=0.0, # forward pass calculations are determinisic
tolerance=0.0, # forward pass calculations are deterministic
fatal=True,
)
rerun_state_machine.validate_result(
result=loss,
rejection_func=torch.isinf,
message="found Inf in local forward loss calculation",
tolerance=0.0, # forward pass calculations are determinisic
tolerance=0.0, # forward pass calculations are deterministic
fatal=True,
)
# Check for spiky loss
Expand All @@ -188,7 +258,7 @@ def loss_func(
context="loss",
),
message="Spiky loss",
tolerance=0.0, # forward pass calculations are determinisic
tolerance=0.0, # forward pass calculations are deterministic
fatal=False,
)

Expand All @@ -208,7 +278,6 @@ def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = Fa

# Get the batch.
timers('batch-generator', log_level=2).start()
global stimer
with stimer(bdata=True):
vp_stage = get_attr_wrapped_model(model, "vp_stage")
(
Expand All @@ -234,7 +303,9 @@ def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = Fa
# Use real (unpadded) cu_seqlens to feed the FLOPs accounting: varlen
# attention only computes work for real tokens within each chunk.
update_seqlen_stats_from_cu_seqlens(cu_seqlens)
cu_seqlens_for_params = cu_seqlens_padded if cu_seqlens_padded is not None else cu_seqlens # TODO(asolergi-nv): Currently there is a bug forcing cu_seqlens to be cu_seqlens_padded
cu_seqlens_for_params = (
cu_seqlens_padded if cu_seqlens_padded is not None else cu_seqlens
) # TODO(asolergi-nv): Currently there is a bug forcing cu_seqlens to be cu_seqlens_padded
packed_seq_params = PackedSeqParams(
qkv_format="thd",
cu_seqlens_q=cu_seqlens_for_params,
Expand All @@ -251,35 +322,45 @@ def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = Fa

with stimer:
if return_schedule_plan:
assert args.overlap_moe_expert_parallel_comm, \
"overlap_moe_expert_parallel_comm must be enabled to return the schedule plan"
assert (
args.overlap_moe_expert_parallel_comm
), "overlap_moe_expert_parallel_comm must be enabled to return the schedule plan"
schedule_plan = model.build_schedule_plan(
tokens, position_ids, attention_mask, labels=labels, loss_mask=loss_mask
)
return schedule_plan, partial(loss_func, loss_mask, model=model)
else:
output_tensor = model(
tokens, position_ids, attention_mask, labels=labels, loss_mask=loss_mask, packed_seq_params=packed_seq_params
tokens,
position_ids,
attention_mask,
labels=labels,
loss_mask=loss_mask,
packed_seq_params=packed_seq_params,
)

# [ModelOpt]: model is needed to access ModelOpt distillation losses
return output_tensor, partial(loss_func, loss_mask, model=model)


def is_dataset_built_on_rank(vp_stage=None, is_packed_sequence=False):
"""Whether the dataset should be built on the current rank."""
args = get_args()
config = core_transformer_config_from_args(args)
if mpu.get_tensor_model_parallel_rank() != 0:
return False
elif is_packed_sequence:
return True
return (
is_first_or_last_pipeline_stage(vp_stage)
or mtp_on_this_rank_func(layout=config.pipeline_model_parallel_layout, mtp_num_layers=config.mtp_num_layers, ignore_virtual=False, vp_stage=vp_stage)
return is_first_or_last_pipeline_stage(vp_stage) or mtp_on_this_rank_func(
layout=config.pipeline_model_parallel_layout,
mtp_num_layers=config.mtp_num_layers,
ignore_virtual=False,
vp_stage=vp_stage,
)


def core_gpt_dataset_config_from_args(args: Any) -> GPTDatasetConfig:
"""Build the GPT (or FIM) dataset config from parsed CLI args."""
tokenizer = build_tokenizer(args)

# Sometimes --data-path is too long, instead we parse it from a file.
Expand Down Expand Up @@ -354,7 +435,6 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None

config = core_gpt_dataset_config_from_args(args)


is_packed_sequence = False
if args.sft:
dataset_type = SFTDataset
Expand All @@ -369,7 +449,9 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None

print_rank_0("> building train, validation, and test datasets for GPT ...")

is_dataset_built = partial(is_dataset_built_on_rank, vp_stage=vp_stage, is_packed_sequence=is_packed_sequence)
is_dataset_built = partial(
is_dataset_built_on_rank, vp_stage=vp_stage, is_packed_sequence=is_packed_sequence
)
train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder(
dataset_type, train_val_test_num_samples, is_dataset_built, config
).build()
Expand Down Expand Up @@ -413,7 +495,8 @@ def get_embedding_ranks(pp_ranks: List[int]):
)
model_cfg = gpt_config_from_args(args)
full_config = pretrain_cfg_container_from_args(args, model_cfg)
pretrain(full_config,
pretrain(
full_config,
train_valid_test_datasets_provider,
partial(model_provider, gpt_builder),
ModelType.encoder_or_decoder,
Expand Down
Loading
Loading