[CarlosGomes98]

Add a method that replaces a parameter and its FP32 master weights without resetting Adam history, for use after loading model weights without optimizer state. It accepts full FP32 values when store_param_remainders is enabled, writing the BF16 high bits with the same rounding convention as AdamFunctorMasterParamRemainder, and supports quantized and DTensor parameters.

The method rejects tensors the optimizer does not own (e.g. param.data), DTensor masters whose mesh, placements or shape differ from the parameter, non-FP32 values and mismatched local shapes, all before mutating any state.

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[timmoon10]

Overall this is a good idea. The optimizer state for master weights can get weird and esoteric (especially the BF16 remainders trick) and it's not reasonable for users to interact with it directly. With some minor tweaks, this API is very straightforward.

[timmoon10]

This is fine functionally, but I think the following signature would be cleaner:

def update_param(self, param: torch.nn.Parameter, values: torch.Tensor) -> None:

The docstring makes you think you're getting rid of param and using master_param for the rest of the training run, when you're actually just updating values in param and in the optimizer state master weights.

[timmoon10]

This is a very Claude-y docstring, which is more showing off its debugging history than providing useful guidance for users.

In fact, it seems we could get across everything important with just:

Suggested change

	"""Replace a parameter and its master weights without resetting Adam history.
	Use this after loading model weights without loading optimizer state. Unlike
	:meth:`set_scaled_state`, this method accepts full FP32 values even when
	``store_param_remainders`` is enabled. It updates the model parameter as well,
	including the high bits needed to reconstruct the master from its remainders.
	Moment buffers are initialized if absent; existing moments and step counters
	are preserved. Distributed callers must synchronize updated parameter shards.
	Call this outside CUDA graph capture. The configured master-weight dtype still
	determines storage precision when remainders are disabled.
	"""Update parameter values and optimizer state for master weights.

I prefer implementing functions with a clear idea of what they are supposed to do, rather than prescribing when you're supposed to call them. Also, store_param_remainders is an internal implementation detail.

[timmoon10]

Do we want to treat get_unscaled_state/set_scaled_state as part of the user-facing API? I imagine it as an internal helper function that expert users like Mcore might use. How much does Mcore depend on it as a stable interface?