Update load mcore checkpoint

[Jintao-Huang]

No description provided.

[gemini-code-assist]

Code Review

This pull request refactors checkpoint loading in Megatron utilities by renaming variables for clarity and ensuring strict=False is consistently applied when loading state dicts for multiple DDP models. In wrap_model, it fixes a return bug when DDP wrapping is disabled, dynamically calculates bucket_size based on num_buckets, disables bucketing for non-first pipeline-parallel ranks, and adds proper CUDA stream synchronization for DDP initialization. The review feedback highlights a potential ZeroDivisionError if num_buckets is set to zero, suggesting a safety check before performing the division.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

If ddp_config.num_buckets is configured to 0, the division num_parameters // ddp_config.num_buckets will raise a ZeroDivisionError. To prevent potential runtime crashes, we should ensure num_buckets is greater than 0 before performing the division. If it is 0 or invalid, it will safely fall back to the default bucket size calculation below.

Suggested change

	if ddp_config.bucket_size is None and getattr(ddp_config, 'num_buckets', None) is not None:
	ddp_config.bucket_size = num_parameters // ddp_config.num_buckets
	if ddp_config.bucket_size is None and getattr(ddp_config, 'num_buckets', None) is not None and ddp_config.num_buckets > 0:
	ddp_config.bucket_size = num_parameters // ddp_config.num_buckets

[ErenAta16]

This reads as a rename in the diff, but it is a correctness fix, and I think that is worth saying out loud because the title and the empty description hide it — it has been open since June 18 with no comments, which is probably why.

state_dict is bound twice in load_mcore_checkpoint, and the second binding shadows the first for everything below it:

410:  state_dict = dist_checkpointing.load_common_state_dict(checkpoint_dir)   # has 'args'
...
446:  sharded_state_dict = _generate_state_dict(args, models, gen_sd_optim,
                                                gen_sd_opt_param_scheduler, gen_sd_rng_state, ...)
467:  state_dict = dist_checkpointing.load(sharded_state_dict, checkpoint_dir, load_strategy)
471:  if 'args' in state_dict and not finetune:
472:      args.consumed_train_samples = getattr(state_dict['args'], 'consumed_train_samples', 0)

At line 471 state_dict is no longer the common state dict. It is whatever dist_checkpointing.load returned for the sharded request built at 446 — model, optimizer, opt_param_scheduler and rng state. args was never asked for, so it is not in the result, 'args' in state_dict is False, and line 472 never runs.

The effect is that args.consumed_train_samples is not restored on resume. There is no exception and no log line, because the missing key is guarded — the run resumes, the counter starts from its default, and the data loader replays samples the run already consumed. Anything keyed off that counter goes with it.

Renaming the first binding to common_state_dict is exactly the right fix: lines 471-472 now read the dict that actually carries args, and the two loads stop sharing a name.

Two suggestions, both about making this land as what it is:

Title and description. "Update load mcore checkpoint" with an empty body is why this sat for six weeks — a reviewer scanning the diff sees nine identical renames and moves on. Something like "fix: restore consumed_train_samples on mcore resume (state_dict shadowing)" plus two lines of description would make it obviously mergeable, and would give anyone bisecting a "resume re-trains on old data" report something to find.

Consider pinning it. A test that resumes from a checkpoint and asserts args.consumed_train_samples is non-zero would stop the shadowing from coming back, since nothing about the current code prevents a third state_dict = ... appearing between 410 and 471 later. If that is too heavy for this path, asserting 'args' in common_state_dict right after line 410 at least makes the assumption explicit.

Line 429's sharded_sd_metadata = state_dict.get('content_metadata') and 430's no_save_optim lookup are both above the reassignment, so they were already reading the right dict — the rename is a no-op for those and the change is safe. I checked each of the nine sites against the two bindings.