core.transformer.moe.megakernel.parameter_bridge

Bridge in-kernel weight-gradient accumulation to MCore DDP parameters.

Module Contents

Functions

dummy_weight_gradient	Return a zero-allocation gradient sentinel that triggers MCore DDP hooks.
main_grad_buffer	Return and validate an optimizer-visible FP32 or BF16 gradient buffer.
finish_weight_gradient	Mark an in-kernel accumulation complete and return a DDP hook-only grad.

API

core.transformer.moe.megakernel.parameter_bridge.dummy_weight_gradient(param: torch.nn.Parameter) → torch.Tensor

Return a zero-allocation gradient sentinel that triggers MCore DDP hooks.

The detached tensor intentionally aliases the parameter storage; callers must never read or mutate its values. This is safe only after the megakernel writes the numerical gradient to main_grad and sets grad_added_to_main_grad=True. MCore’s DDP hook then clears .grad without reading or accumulating this sentinel.

core.transformer.moe.megakernel.parameter_bridge.main_grad_buffer(param: torch.nn.Parameter) → torch.Tensor

Return and validate an optimizer-visible FP32 or BF16 gradient buffer.

core.transformer.moe.megakernel.parameter_bridge.finish_weight_gradient(param: torch.nn.Parameter) → torch.Tensor

Mark an in-kernel accumulation complete and return a DDP hook-only grad.