[optim]--rematerialize-param-from-master-weight: save the bf16 weight backup in colocate

[yueming-yuan]

Stacked on #1571.

TLDR

On GLM 5.2, 64 * GB300, before -> after optimization:
CPU memory 858.2 GB -> 729.2 GB (-129 GB)
weight update time 76.3 s -> 73.8 s (-2.5 s)
wake_up time (include rematerialize) 4.6 s -> 4.8 s (+0.2 s)

Motivation

In colocate, the weights_backuper keeps a pinned CPU copy of the bf16 weights (2 bytes/param/rank) to (a) restore params after the TMS rollout pause and (b) feed update_weights while the actor is paused. On GLM-5.2 744B that copy is 26-37.5 GB/rank — 104-150 GB per node sitting inside a rollout-peak budget that measured 858/898 GiB.

Both jobs can be done without it:

• update_weights runs with the param buffer still resident (phase 1 of the sleep pauses grads/main/mv only) and reads live GPU weights; the buffer is dropped without backup right after the update. Engine pools are sized on an empty GPU at startup, so engine + 2P coexist only during the short update window — the KV pool is untouched.
• restore at wake replays exactly what produced the bf16 weights at step end: `_copy_main_params_to_model_params` (fp32 main -> bf16 cast) + param all-gather. Bit-identical by construction.
• Only tensors not derivable from fp32 mains keep a KB-scale pinned backup: expert_bias buffers and fp32-dtype params (their optimizer "main" aliases the param itself). update_weights reads those from the backup since their TMS region is paused.

Guarantees

• `--check-rematerialize-param-from-master-weight N` (default 2): per-tensor SHA256 recorded at backup, re-verified after restore, for the first N cycles. N=2 also catches param-gather state-machine damage, which surfaces in the following step's weights.
• Init-time coverage assert: every named parameter must be in the DDP param buffers (restorable via cast+AG) or in the extras backup — fails loudly on e.g. frozen params. (DDP buffer membership is the correct local criterion: all distributed-optimizer structures only cover this rank's owned shard under DP>1.)
• 8 config asserts exclude ref/teacher/old-actor tags, LoRA, precision-aware optimizers, overlap-param-gather, and --disable-compute-advantages-and-returns (the per-cycle restore runs in that block). critic-only warmup steps are fenced (they rerun update_weights without an intervening actor wake_up; Skip redundant actor weight update during critic-only steps #1588 removes the redundant updates and lifts the fence). --debug-train-only silently disables the flag and the critic actor keeps the baseline pause path - neither ever runs update_weights, where the param buffer is paused.
• update_weights' weight source asserts every tensor is either in the extras backup or in the DDP param buffers - a new buffer admitted into the sync source fails loudly instead of handing out a pointer into the paused region.

Measured

	before	after
host rollout peak, max node (744B, 16x GB300)	858.2 GB	729.2 GB (-129)
engine KV pool	1.64M tokens	unchanged
steady step time	~1000 s	972-989 s (restore cast+AG ~ +1.5 s, per-step backup D2H removed)
weight bit-exactness	pinned-copy restore	SHA256-identical x3 cycles @744B TP8xPP4xDP2xEP16, x4 @Qwen3-4B TP2xDP2

[gemini-code-assist]

Code Review

This pull request introduces the --restore-weights-from-fp32-main optimization, which reconstructs low-precision model weights from the distributed optimizer's fp32 main parameters instead of maintaining a full pinned CPU copy during rollout pauses, thereby reducing host memory usage. Feedback on these changes highlights several robustness improvements: first, frozen parameters (requires_grad=False) should be included in the extras backup to prevent startup assertion failures; second, direct attribute access on model modules for buffers should use getattr to avoid potential AttributeErrors; and third, defensive checks should be added before invoking optimizer and model chunk synchronization methods to handle edge cases like disabled optimizers.

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]

Parameters that do not require gradients (requires_grad=False, i.e., frozen parameters) are not managed by the distributed optimizer or DDP buffers. If a model contains any frozen parameters, they will be missed by the restore path and cause a startup assertion failure in _assert_restore_coverage. Including them in named_restore_extras ensures they are correctly backed up and restored.

Suggested change

	def named_restore_extras(model: Sequence[torch.nn.Module]) -> Iterator[tuple[str, torch.Tensor]]:
	"""Tensors NOT reconstructable from the distributed optimizer's fp32 main
	params: expert_bias buffers (the only buffers synced today, see
	_named_params_and_buffers_vanilla) and fp32-dtype parameters, whose
	optimizer "main" is a view of the param itself (shard_fp32_groups)."""
	for vp_stage, model_module in enumerate(model):
	for name, buffer in model_module.named_buffers():
	if "expert_bias" in name:
	yield f"vp_stages.{vp_stage}.{strip_param_name_prefix(name)}", buffer
	for name, param in model_module.named_parameters():
	if param.dtype == torch.float32:
	yield f"vp_stages.{vp_stage}.{strip_param_name_prefix(name)}", param
	def named_restore_extras(model: Sequence[torch.nn.Module]) -> Iterator[tuple[str, torch.Tensor]]:
	"""Tensors NOT reconstructable from the distributed optimizer's fp32 main
	params: expert_bias buffers (the only buffers synced today, see
	_named_params_and_buffers_vanilla), fp32-dtype parameters, and parameters
	that do not require gradients (frozen parameters)."""
	for vp_stage, model_module in enumerate(model):
	for name, buffer in model_module.named_buffers():
	if "expert_bias" in name:
	yield f"vp_stages.{vp_stage}.{strip_param_name_prefix(name)}", buffer
	for name, param in model_module.named_parameters():
	if param.dtype == torch.float32 or not param.requires_grad:
	yield f"vp_stages.{vp_stage}.{strip_param_name_prefix(name)}", param

[gemini-code-assist]

Using direct attribute access like model_module.buffers and model_module.expert_parallel_buffers can raise an AttributeError if the model chunks are not wrapped in Megatron's DDP wrapper or if the Megatron version/configuration differs (e.g., if expert parallel buffers are not initialized). Using getattr with a default empty list makes this check robust and prevents potential crashes.

Suggested change

	for model_module in self.model:
	for buffer in model_module.buffers + model_module.expert_parallel_buffers:
	restorable.update(id(p) for p in buffer.params)
	for model_module in self.model:
	buffers = getattr(model_module, "buffers", []) + getattr(model_module, "expert_parallel_buffers", [])
	for buffer in buffers:
	restorable.update(id(p) for p in getattr(buffer, "params", []))

[gemini-code-assist]

To adhere to defensive programming guidelines, we should add safety guards before calling _copy_main_params_to_model_params and start_param_sync. For example, if --debug-disable-optimizer is enabled, self._ctx.optimizer will be None, leading to an AttributeError when trying to call _copy_main_params_to_model_params. Similarly, checking if model_chunk has start_param_sync prevents crashes if any model chunks are not wrapped in DDP.

Suggested change

	optimizers = getattr(self._ctx.optimizer, "chained_optimizers", [self._ctx.optimizer])
	for optimizer in optimizers:
	optimizer._copy_main_params_to_model_params()
	for model_chunk in self._ctx.model_chunks:
	model_chunk.start_param_sync(force_sync=True)
	optimizers = getattr(self._ctx.optimizer, "chained_optimizers", [self._ctx.optimizer])
	for optimizer in optimizers:
	if optimizer is not None and hasattr(optimizer, "_copy_main_params_to_model_params"):
	optimizer._copy_main_params_to_model_params()
	for model_chunk in self._ctx.model_chunks:
	if hasattr(model_chunk, "start_param_sync"):
	model_chunk.start_param_sync(force_sync=True)