API Reference¶
TL;DR
DMuon exposes five surface areas: setup (dedicate_params, install_patch),
optimizer (Muon, NewtonSchulz, NS functions and constants), state
management (no_sync, wait_all_reduces, replicate-broadcast helpers,
DedicatedCommContext), diagnostics (summarize_param_groups,
summarize_comm_plan, collect_forward_unshard_profile), and checkpointing
(get/set_model/optimizer_state_dict).
Start with dedicate_params + Muon; reach for the rest when you need fine-grained
control.
Setup¶
dedicate_params¶
Called once before fully_shard(). Assigns each Muon-target parameter to
a single owner rank and registers the per-layer forward/backward hooks. See
Custom Hook Boundaries and
Z2 vs Z3 Modes for the two most common
customization points.
By default, dedicate_params() uses process_group_policy="isolated". DMuon
clones the DP/HSDP/TP mesh process groups for its own collectives, so external
trainer logging, metrics, and checkpoint collectives cannot enter the same
NCCL communicator sequence as DMuon's async post-step publish. Pass
process_group_policy="shared" only when you explicitly want the historical
behavior of reusing the caller-provided DeviceMesh groups.
isolated only changes process-group ownership. It does not add a step-end
barrier by default, so DMuon's async publish can still overlap with the next
training step. For debugging suspected process-group ordering issues, set
DMUON_ISOLATED_PG_BARRIER=1 to force a strict post-step fence over DMuon-owned
groups. Do not enable that fence for performance measurements unless the
experiment is explicitly measuring the no-overlap behavior.
dedicate_params ¶
dedicate_params(model: Module, mesh: DeviceMesh, predicate: Callable[[str, Parameter], bool], compute_dtype: dtype = None, reshard_after_forward: bool = True, replicate_mesh: Optional[DeviceMesh] = None, hook_boundary_predicate: Optional[Callable[[Module], bool]] = None, hook_boundary_strict: bool = True, hook_boundary_resolver: Optional[HookBoundaryResolver] = None, assignment_group_key_fn: Optional[AssignmentGroupKeyFn] = None, route_hint_fn: Optional[RouteHintFn] = None, max_owners_per_group: Optional[int] = None, owner_strategy: Optional[str] = None, owner_cost_model: Optional[str] = None, hsdp_column_balance: Optional[bool] = None, pack_small_params: Optional[bool] = None, tp_buffer_reuse: bool | str = False, replicate_broadcast_bucket_mb: float = 0.0, muon_forward_unshard: str = 'broadcast', delay_stage2_to_optimizer: bool = True, param_policy: Optional[Mapping[str, object]] = None, param_policy_fn: Optional[ParamPolicyFn] = None, process_group_policy: Optional[str] = None) -> dict[Parameter, int]
Mark parameters for dedicated ownership and register communication hooks.
Parameters satisfying predicate are assigned to owner ranks via a
balanced partition algorithm. Each marked parameter will be automatically
ignored by subsequent fully_shard() calls (requires the monkey-patch
installed automatically by :mod:dmuon).
Communication hooks are registered at the layer level (e.g., on
model.layers[i]), not on individual sub-modules. This minimizes
CPU launch overhead by batching all broadcasts/reduces per layer.
Tensor parallelism is supported transparently: if a parameter is a
DTensor sharded on a mesh dim that is NOT named in mesh or
replicate_mesh (i.e. the TP axis), DMuon will — at optimizer step
time — gather the full matrix at a designated TP owner, run
Newton-Schulz locally, and scatter the per-shard update back to
the TP group. No TP mesh argument is required; simply pass the DP
slice of your 3D mesh as mesh / replicate_mesh, matching the
FSDP2 convention (fully_shard(mesh=mesh["replicate","shard"])).
Typical 3D (replicate × shard × tp) call order::
mesh3d = init_device_mesh("cuda", (R, G, T),
mesh_dim_names=("replicate","shard","tp"))
parallelize_module(model, mesh3d["tp"], {...}) # TP first
dmuon.dedicate_params( # DMuon BEFORE fully_shard
model,
mesh=mesh3d["shard"],
replicate_mesh=mesh3d["replicate"],
predicate=...,
)
fully_shard(model, mesh=mesh3d["replicate","shard"])
Args:
model: The model whose parameters to partition.
mesh: 1D DeviceMesh over the shard dimension (a.k.a. dp_group).
When replicate_mesh is provided, this becomes the shard axis
of the HSDP 2D mesh. Must be a named sub-mesh (constructed via
init_device_mesh(..., mesh_dim_names=...)) whenever TP is
present so the TP axis can be inferred by name-set difference.
predicate: Callable (param_name, param) -> bool. Parameters
returning True will use dedicated ownership.
compute_dtype: Optional dtype for communication (e.g., torch.bfloat16).
reshard_after_forward: If True (default), reshard dedicated params after
forward and re-broadcast in backward. If False (SHARD_GRAD_OP mode),
keep params unsharded through forward+backward, eliminating backward
broadcasts at the cost of higher memory.
replicate_mesh: Optional 1D DeviceMesh over the replicate dimension.
When provided, dedicate_params accepts a HSDP-style 2D layout; the
LPT partition then balances globally over G·R owner slots
(G = shard size, R = replicate size).
hook_boundary_predicate: Optional (module) -> bool selector for the
hook attachment module. When set, DMuon registers its pre/post
forward hooks on the lowest ancestor of each dedicated param
where the predicate is True. Use this to align hook boundaries
with your FSDP2 fully_shard boundaries — e.g. treat the whole
ViT as one hook site even though its parameters are distributed
across ranks. Leave as None to use the built-in layers.N /
blocks.N heuristic (:func:_find_layer_module).
hook_boundary_strict: When True (default) and hook_boundary_predicate
is given, raise if a dedicated param has no ancestor matching the
predicate. When False, fall back to the param's direct parent
module. Strict is recommended to avoid silent per-sub-module hooks.
hook_boundary_resolver: Optional (param_fqn, param) -> module hook
selector for models whose execution boundary is not an ancestor of
the parameter. It takes precedence over hook_boundary_predicate.
assignment_group_key_fn: Optional (param_fqn, param) -> key selector
used by the owner assignment pass. Use this to align owner packing
with communication boundaries.
route_hint_fn: Optional (param_fqn, param) -> route selector used
before parameters are replaced by DMuon placeholders. This is
required for DMuon-managed sharded base AdamW parameters because
their initial local shard must be captured while the full tensor
is still present on the rank. Supported route hints are
"muon", "adamw", and "sharded_adamw".
max_owners_per_group: Optional cap on distinct owner slots used by one
assignment group.
owner_strategy: DP/HSDP owner assignment strategy. "lpt" is the
production default; "round_robin" and "rank0" are intended
for benchmark ablations. When omitted, DMUON_OWNER_STRATEGY can
override the default "lpt" for benchmark jobs.
owner_cost_model: LPT cost model. "optimizer" uses shape-aware
matrix optimizer cost plus footprint. "numel" is a diagnostic
numel-only ablation. When omitted, DMUON_OWNER_COST_MODEL can
override the default "optimizer" for benchmark jobs.
hsdp_column_balance: Whether HSDP LPT should balance shard-column load
before per-owner load. Disable only for placement ablations. When
omitted, DMUON_HSDP_COLUMN_BALANCE can override the default
True.
pack_small_params: Whether same-layer small parameters are merged into
packed allocation units before owner assignment. Disable only for
true per-matrix baseline jobs. When omitted,
DMUON_OWNER_PACK_SMALL_PARAMS can override the default True.
tp_buffer_reuse: Optional TP gather/scatter scratch-buffer reuse policy.
Accepts False/True or "gather", "scatter", "all".
replicate_broadcast_bucket_mb: Optional HSDP post-step publish bucket
size in MiB. 0 keeps one coalesced publish per hook group.
muon_forward_unshard: Forward unshard placement for Muon-routed
parameters. "broadcast" keeps the existing owner-to-all
publish path. "all_gather" keeps the owner-side full-matrix
update, then scatters the updated tensor into rank-local shards
after the optimizer step and reconstructs the forward view with an
FSDP-style all-gather. The all-gather mode is experimental and is
currently supported only for non-TP Muon parameters.
delay_stage2_to_optimizer: When True, backward only waits the Stage-1
shard reduce for buffer lifetime; HSDP Stage-2 replicate reduce is
waited by per-group optimizer preparation. This is the default
because it preserves overlap between late reduce tails and earlier
optimizer/publish work.
param_policy: Optional structured per-parameter policy. The policy starts
from defaults and applies ordered overrides whose
name tokens match the full parameter name. Supported fields
include route, param_dtype, grad_dtype, output_dtype,
cast_forward_inputs, master_dtype, optim_dtype, and
muon_forward_unshard. Prefer this over route_hint_fn for pure
DMuon integrations that need mixed precision policy.
param_policy_fn: Optional (param_fqn, param) -> policy callable for
dynamic policy decisions. Mutually exclusive with param_policy.
process_group_policy: DMuon process group ownership policy.
"isolated" (default, or DMUON_PROCESS_GROUP_POLICY) clones
DMuon's DP/HSDP/TP process groups so external trainer collectives
cannot interleave with DMuon async communication on the same NCCL
communicator. "shared" preserves historical behavior and uses
the caller-provided mesh groups directly.
Returns:
Assignment dict mapping each dedicated parameter to its DP owner
coord (int in 1D shard-only mode, Tuple[int, int] in HSDP
mode). TP owner selection is not returned — it is intrinsic
to each DTensor's TP group and resolved at hook-registration time
(DedicatedParam.is_tp_owner).
dedicate_params_ddp¶
DDP-path setup for dedicated parameters when the data-parallel model is replicated instead of FSDP2-sharded.
dedicate_params_ddp ¶
dedicate_params_ddp(model: Module, mesh: DeviceMesh, predicate: Callable[[str, Parameter], bool], compute_dtype: Optional[dtype] = None, hook_boundary_predicate: Optional[Callable[[Module], bool]] = None, hook_boundary_strict: bool = True, hook_boundary_resolver: Optional[HookBoundaryResolver] = None, assignment_group_key_fn: Optional[AssignmentGroupKeyFn] = None, max_owners_per_group: Optional[int] = None, owner_strategy: Optional[str] = None, owner_cost_model: Optional[str] = None, pack_small_params: Optional[bool] = None, param_policy: Optional[Mapping[str, object]] = None, param_policy_fn: Optional[ParamPolicyFn] = None, process_group_policy: Optional[str] = None) -> dict[Parameter, int]
DDP-path variant of :func:dedicate_params.
Every rank keeps the full parameter live on the module. Ownership
applies only to (1) who runs Newton-Schulz, (2) the dist.reduce
destination after backward, and (3) the dist.broadcast source
after optim.step.
Compared with :func:dedicate_params, this entry:
- Accepts a 1D mesh only. 2D mesh (HSDP) raises — HSDP users
should continue to use
dedicate_params+fully_shard. - Does not replace dedicated parameters with 0-size placeholders
on non-owner ranks; the original
nn.Parameterstays live. - Does not register forward-broadcast / reshard hooks.
DedicatedStatestill attaches hooks at the layer boundary, butunshard/resharddegrade to no-op on DDP groups; the real work isreduce_gradsafter backward +post_step_broadcastafteroptim.step.
Args:
model: Model to partition.
mesh: 1D DeviceMesh over the data-parallel world.
predicate: Callable (param_name, param) -> bool selecting
dedicated parameters.
compute_dtype: Optional dtype for communication.
hook_boundary_predicate: Same semantics as in
:func:dedicate_params. Default heuristic uses
layers.N / blocks.N.
hook_boundary_strict: Same as :func:dedicate_params.
hook_boundary_resolver: Same as :func:dedicate_params.
assignment_group_key_fn: Same as :func:dedicate_params.
max_owners_per_group: Same as :func:dedicate_params.
owner_strategy: Same as :func:dedicate_params.
owner_cost_model: Same as :func:dedicate_params.
pack_small_params: Same as :func:dedicate_params.
param_policy: Same structured policy surface as :func:dedicate_params.
param_policy_fn: Same dynamic policy surface as :func:dedicate_params.
process_group_policy: Same process-group ownership policy as
:func:dedicate_params.
Returns:
Assignment dict mapping each dedicated parameter to its owner
rank (int — 1D shard-only).
dedicate_params_ddp_tp¶
DDP-path setup for dedicated parameters when tensor parallelism is active inside each replicated data-parallel group.
dedicate_params_ddp_tp ¶
dedicate_params_ddp_tp(model: Module, mesh: DeviceMesh, predicate: Callable[[str, Parameter], bool], compute_dtype: Optional[dtype] = None, hook_boundary_predicate: Optional[Callable[[Module], bool]] = None, hook_boundary_strict: bool = True, hook_boundary_resolver: Optional[HookBoundaryResolver] = None, assignment_group_key_fn: Optional[AssignmentGroupKeyFn] = None, max_owners_per_group: Optional[int] = None, owner_strategy: Optional[str] = None, owner_cost_model: Optional[str] = None, pack_small_params: Optional[bool] = None, tp_buffer_reuse: bool | str = False, param_policy: Optional[Mapping[str, object]] = None, param_policy_fn: Optional[ParamPolicyFn] = None, process_group_policy: Optional[str] = None) -> dict[Parameter, int]
DDP-path variant that allows TP-sharded DTensor dedicated params.
This API is intentionally separate from :func:dedicate_params_ddp so
pure DDP keeps rejecting DTensor parameters by default. It is the
supported DDP + TP path:
- Apply tensor parallelism first.
- Call
dedicate_params_ddp_tp(model, mesh["dp"], ...). - Call
replicate_tp(model, mesh["dp"])for non-dedicated params.
Each TP-sharded dedicated parameter still has a DP owner. All TP ranks in that DP-owner replica gather their local reduced grads to a TP owner for full-matrix Newton-Schulz, scatter update shards back, then DDP broadcasts those updated shards across the DP mesh.
process_group_policy has the same semantics as :func:dedicate_params:
"isolated" (default) makes DMuon clone its DP/TP process groups, while
"shared" reuses the caller-provided mesh groups.
replicate¶
DDP-style replication helper for non-dedicated parameters.
replicate ¶
DDP-style replication for non-dedicated parameters.
Companion to :func:dmuon.dedicate_params_ddp. dedicate_params_ddp
handles parameters selected by the predicate (*.proj.weight, ...);
everything else (LayerNorm, head, embeddings, ...) is handled
here by attaching a post-accumulate-grad hook that averages gradients
across the data-parallel world at end-of-backward.
The design mirrors FSDP2's post-backward pattern and DMuon's own
DedicatedState._queue_root_post_backward_callback (state.py):
per-param hooks enqueue the param, and the autograd root callback
flushes the whole batch in one dist._coalescing_manager so NCCL
fuses the per-bucket all-reduces into a single kernel.
replicate is intentionally separate from FSDP2's
fully_shard. The two are per-parameter mutually exclusive; calling
replicate on a model where FSDP2 already manages the same param
raises. This is the cleanest boundary — users pick one path per
non-dedicated param.
2D mesh (HSDP-minimal — every rank keeps a full replica of non-dedicated
params and averages across the 2D world) is explicitly NOT supported in
P1. Use fully_shard(mesh=hsdp) for HSDP; use replicate(mesh=1D)
for pure DDP.
replicate ¶
Install DDP-style replication for non-dedicated parameters.
Must be called after :func:dedicate_params_ddp (or
:func:dedicate_params; though in the FSDP2 path users normally
call fully_shard instead). Dedicated parameters are skipped
automatically.
Example::
mesh = init_device_mesh("cuda", (world_size,))
dmuon.dedicate_params_ddp(model, mesh, predicate=...)
dmuon.replicate(model, mesh=mesh)
optimizer = dmuon.Muon(model, lr=..., adamw_lr=...)
The returned object is the input model itself (identity) — no
wrapping, no .module unwrap needed downstream. This mirrors
:func:torch.distributed.fsdp.fully_shard UX.
Args: model: The model to attach replicated-group state to. mesh: 1D DeviceMesh for the data-parallel world.
Raises:
ValueError: If mesh is not 1D.
RuntimeError: If replicate was already called on model,
or if a managed parameter is already under FSDP2.
Returns:
model (for call chaining).
replicate_tp ¶
Install DDP-style replication for TP-only DTensor parameters.
This is the TP-aware companion to :func:dmuon.dedicate_params_ddp_tp.
It keeps the conservative FSDP-DTensor rejection in :func:replicate
intact while allowing tensor-parallel DTensors to all-reduce their local
gradient shards over the data-parallel mesh.
get_replicated_group ¶
Return the _ReplicatedGroup attached to model, or None.
replicate_tp¶
TP-aware companion to dedicate_params_ddp_tp() for non-dedicated
DTensor parameters.
replicate_tp ¶
Install DDP-style replication for TP-only DTensor parameters.
This is the TP-aware companion to :func:dmuon.dedicate_params_ddp_tp.
It keeps the conservative FSDP-DTensor rejection in :func:replicate
intact while allowing tensor-parallel DTensors to all-reduce their local
gradient shards over the data-parallel mesh.
install_patch¶
import dmuon calls this automatically. You should never need to call it
directly unless you are constructing a DMuon environment without the normal
import path.
install_patch ¶
Install the FSDP2 monkey-patch that makes fully_shard skip
dedicated parameters.
Called automatically on import dmuon — users do not normally
invoke this directly. After patching, any subsequent call to
fully_shard() filters out parameters previously marked with
_dedicated_owner_rank by :func:dedicate_params, leaving them
under DMuon's ownership instead of FSDP2's uniform sharding.
Safe to call repeatedly (idempotent). The reverse operation is
:func:uninstall_patch.
Patched function: torch.distributed.fsdp._fully_shard._fsdp_init.
_get_managed_states.
Optimizer¶
Muon¶
The primary optimizer class. Manages Muon (Newton-Schulz + momentum) on
matrix-routed dedicated parameters and AdamW on the base path in a single
object. The base path can be ordinary FSDP2-managed parameters or
DMuon-managed sharded AdamW parameters selected via param_policy or the
legacy route_hint_fn.
Compatible with torch.optim.lr_scheduler.
Muon ¶
Muon(model: Module, lr: float = 0.02, momentum: float = 0.95, weight_decay: float = 0.0, ns_steps: int = 5, adamw_lr: float = 0.001, adamw_betas: tuple[float, float] = (0.9, 0.999), adamw_weight_decay: float = 0.01, adamw_eps: float = 1e-08, param_groups: Optional[list[dict]] = None, ns_backend: Union[str, NewtonSchulz] = 'gram', nesterov: bool = True, replicate_async: bool = True, record_step_profile: bool = False, group_prepare_ahead: bool = True, tp_distributed_gram: bool = False, tp_distributed_gram_policy: str = 'beneficial', tp_distributed_gram_max_factor_to_scatter_ratio: float = 0.5, post_step_prefetch_groups: int = 0, post_step_prefetch_sharded_adamw: bool = False, sharded_adamw_unshard_separate_stream: bool = False, forward_prefetch_depth: int = 1, first_step_progress_log: bool = True, first_step_progress_log_limit: int = 8)
Bases: Optimizer
Muon optimizer for DMuon distributed training.
Manages two types of parameters: - Dedicated params (proj layers): Newton-Schulz orthogonalization with momentum. Only the owner rank computes the update. - Symmetric params (layernorm, embedding): AdamW, updated by all ranks on their FSDP2 shards.
Args:
model: Model with dedicate_params and fully_shard already applied.
lr: Muon learning rate for dedicated params.
momentum: Momentum coefficient for dedicated params.
weight_decay: Weight decay for dedicated params.
ns_steps: Number of Newton-Schulz iterations.
adamw_lr: AdamW learning rate for symmetric params.
adamw_betas: AdamW beta coefficients.
adamw_weight_decay: AdamW weight decay.
adamw_eps: AdamW epsilon.
param_groups: Optional PyTorch-style semantic parameter groups. Each
user group is lowered into a Muon subgroup and an AdamW subgroup
so schedulers and checkpoint metadata can keep per-group
hyperparameters without exposing DMuon internals.
ns_backend: Newton-Schulz backend configuration. Accepts a string
shorthand ("gram" or "direct") or a fully configured
:class:~dmuon.NewtonSchulz object for custom coefficients::
# String shorthand (default coefficients)
optimizer = dmuon.Muon(model, ns_backend="gram")
# Custom coefficients
ns = dmuon.NewtonSchulz("direct", coefficients=dmuon.YOU_COEFFICIENTS)
optimizer = dmuon.Muon(model, ns_backend=ns)
``"gram"`` uses Gram-space NS with SYRK acceleration and restarts.
``"direct"`` uses classic parameter-space NS (Muon/Moonlight).
NS always runs on the full (un-sharded) matrix: for TP-sharded
parameters the runtime reassembles the matrix via All-to-All
before calling NS (see ``tp_design.md``).
nesterov: If True (default), use Nesterov momentum lookahead
before NS orthogonalization: ``ns_input = grad + μ * buf``.
Recommended by original Muon paper and used by Moonlight.
replicate_async: If True (default), publish owner updates asynchronously
and consume the events in the next forward. If False, drain the
publish path inside ``step()`` for deterministic timing. When TP
dedicated parameters are present, DMuon currently keeps this path
synchronous for correctness.
record_step_profile: If True, record CUDA-event timing for optimizer
phases and expose it via ``consume_last_step_profile()``.
group_prepare_ahead: If True, prepare the next group's reduced grads
while the current group's optimizer math runs.
tp_distributed_gram: Enable the TP-aware distributed Gram path for
TP-sharded matrices.
tp_distributed_gram_policy: Policy for the distributed Gram path;
``"beneficial"`` only uses it when the factor payload is expected
to be smaller than scattering the full update.
tp_distributed_gram_max_factor_to_scatter_ratio: Maximum factor-payload
to full-scatter byte ratio allowed by the ``"beneficial"`` policy.
first_step_progress_log: Print a small number of first optimizer step
progress messages on owner ranks. This makes one-time NS/SYRK
per-parameter dispatch or autotune visible in cluster logs.
first_step_progress_log_limit: Maximum number of per-parameter first-step
progress messages to print per rank.
Example::
import dmuon
from torch.distributed.fsdp import fully_shard
dmuon.dedicate_params(model, mesh, predicate=lambda n, p: "proj" in n)
for layer in model.layers:
fully_shard(layer, mesh=mesh)
fully_shard(model, mesh=mesh)
optimizer = dmuon.Muon(model, lr=0.02, momentum=0.95)
for batch in dataloader:
optimizer.zero_grad()
loss = model(batch).loss
loss.backward()
optimizer.step()
last_muon_grad_clip_stats
property
¶
last_muon_grad_clip_stats: Optional[MuonGradClipStats]
Stats from the most recent Muon gradient clipping call.
consume_last_step_profile ¶
Return the last step's CUDA event timings, synchronizing if needed.
clip_grad_norm_ ¶
clip_grad_norm_(max_norm: Optional[float], *, norm_type: Optional[float] = None, error_if_nonfinite: Optional[bool] = None, foreach: Optional[bool] = None, strategy: Optional[object] = None) -> MuonGradClipStats
Clip DMuon dedicated/Muon gradients only.
Ordinary AdamW parameters are intentionally excluded. Use
torch.nn.utils.clip_grad_norm_ for those param.grad tensors and
call this method as the DMuon-specific extra line.
step ¶
Perform a single optimization step.
Internally: 1. Prepares Muon gradients: wait reduce tails and gather TP shards into full gradients on TP owners. 2. Runs Muon (momentum + NS + update) on owned dedicated params. 3. Runs AdamW on FSDP2 symmetric params. 4. Publishes updated dedicated params. In sync mode this happens as one post-step phase; in async mode each group dispatches its scatter/broadcast immediately after that group's Muon update, and the next forward consumes the pending event before reading it.
zero_grad ¶
Clear gradients.
Clears FSDP2 params' gradients and dedicated params' accumulated gradients (from gradient accumulation). Dedicated params' _reduced_grad is normally cleared in step(), but is also cleared here for safety.
NewtonSchulz¶
Configurable NS backend object. Pass to Muon(ns_backend=...) to select the
algorithm variant, override coefficients, or enable deterministic mode. See
Newton-Schulz Variants for a full comparison.
NewtonSchulz ¶
NewtonSchulz(backend: str = 'gram', kernel: str = 'auto', coefficients: Optional[list[list[float]]] = None, restart_iterations: Optional[list[int]] = None, deterministic: bool = False)
Configurable Newton-Schulz backend.
Encapsulates the algorithm variant, coefficients, and SYRK kernel
backend so they can be passed as a single object to :class:~dmuon.Muon.
Args:
backend: "gram" (default) for Gram-space NS with SYRK
acceleration and restarts, or "direct" for classic
parameter-space NS (Muon/Moonlight formulation).
kernel: SYRK kernel backend to use inside Gram NS:
* ``"auto"`` (default) — pick the best available on this GPU.
SM80/87 → ``cute_sm80``, SM90+ with quack installed →
``quack``, otherwise ``cublas``.
* ``"quack"`` — Tri Dao quack SYRK (SM90+, soft dep). Raises
at construction if unavailable.
* ``"cute_sm80"`` — DMuon-internal CuteDSL kernel (SM80/87).
* ``"cublas"`` — universal fallback; bit-exact across runs.
Env var override: ``DMUON_NS_KERNEL`` takes precedence only
when this argument is left at ``"auto"``.
coefficients: Per-step ``(a, b, c)`` coefficients. ``None``
uses :data:`POLAR_EXPRESS_COEFFICIENTS`.
restart_iterations: Restart positions for Gram-space NS.
``None`` uses ``[2]``. Ignored when *backend* is
``"direct"``.
deterministic: Back-compat alias for ``kernel="cublas"``. When
``True`` and ``kernel`` is still ``"auto"``, the kernel is
forced to ``cublas`` (bit-exact reproducibility). Explicit
``kernel=`` wins over ``deterministic`` — if both are given
and they disagree, a warning is emitted and ``kernel`` wins.
Example::
import dmuon
# Default (Gram-space, auto-selected kernel)
ns = dmuon.NewtonSchulz()
# Force cuBLAS (reproducible across runs)
ns = dmuon.NewtonSchulz(kernel="cublas")
ns = dmuon.NewtonSchulz(deterministic=True) # equivalent
# SM90+ explicit quack
ns = dmuon.NewtonSchulz(kernel="quack")
optimizer = dmuon.Muon(model, lr=0.02, ns_backend=ns)
local ¶
Run NS on a full (un-sharded) matrix.
The runtime guarantees the matrix handed in here is the full
logical gradient: for pure-DP params the owner already holds
the full tensor; for TP-sharded params the TP gather
step (dmuon._backends.fsdp2.group.tp_gather_grads) has
reassembled the full matrix on the TP owner before this call.
newton_schulz¶
Standalone NS function (Gram-space by default). Use directly for NS outside the optimizer loop.
newton_schulz ¶
Newton-Schulz orthogonalization with tiered hardware backends.
Backend selection (auto-detected at import time):
- SM80+ & CuteDSL SYRK: lower-triangle + mirror-write kernel (50% tile savings)
- Fallback: @torch.compile pure PyTorch
Two NS modes:
- :func:newton_schulz: Gram-space NS on a full matrix (default public entry).
- :func:direct_newton_schulz: classic parameter-space NS.
TP support lives entirely in the runtime layer (dmuon._backends.fsdp2):
for TP-sharded parameters the runtime does a TP gather so the TP owner
sees the full matrix, then calls one of the functions above — the
NS algorithms themselves are TP-agnostic. See docs/guides/tp-support.md.
Gram NS iteration logic is adapted from Dao-AILab/gram-newton-schulz, including per-step coefficients, restart mechanism, and mixed-precision pipeline.
NewtonSchulz ¶
NewtonSchulz(backend: str = 'gram', kernel: str = 'auto', coefficients: Optional[list[list[float]]] = None, restart_iterations: Optional[list[int]] = None, deterministic: bool = False)
Configurable Newton-Schulz backend.
Encapsulates the algorithm variant, coefficients, and SYRK kernel
backend so they can be passed as a single object to :class:~dmuon.Muon.
Args:
backend: "gram" (default) for Gram-space NS with SYRK
acceleration and restarts, or "direct" for classic
parameter-space NS (Muon/Moonlight formulation).
kernel: SYRK kernel backend to use inside Gram NS:
* ``"auto"`` (default) — pick the best available on this GPU.
SM80/87 → ``cute_sm80``, SM90+ with quack installed →
``quack``, otherwise ``cublas``.
* ``"quack"`` — Tri Dao quack SYRK (SM90+, soft dep). Raises
at construction if unavailable.
* ``"cute_sm80"`` — DMuon-internal CuteDSL kernel (SM80/87).
* ``"cublas"`` — universal fallback; bit-exact across runs.
Env var override: ``DMUON_NS_KERNEL`` takes precedence only
when this argument is left at ``"auto"``.
coefficients: Per-step ``(a, b, c)`` coefficients. ``None``
uses :data:`POLAR_EXPRESS_COEFFICIENTS`.
restart_iterations: Restart positions for Gram-space NS.
``None`` uses ``[2]``. Ignored when *backend* is
``"direct"``.
deterministic: Back-compat alias for ``kernel="cublas"``. When
``True`` and ``kernel`` is still ``"auto"``, the kernel is
forced to ``cublas`` (bit-exact reproducibility). Explicit
``kernel=`` wins over ``deterministic`` — if both are given
and they disagree, a warning is emitted and ``kernel`` wins.
Example::
import dmuon
# Default (Gram-space, auto-selected kernel)
ns = dmuon.NewtonSchulz()
# Force cuBLAS (reproducible across runs)
ns = dmuon.NewtonSchulz(kernel="cublas")
ns = dmuon.NewtonSchulz(deterministic=True) # equivalent
# SM90+ explicit quack
ns = dmuon.NewtonSchulz(kernel="quack")
optimizer = dmuon.Muon(model, lr=0.02, ns_backend=ns)
local ¶
Run NS on a full (un-sharded) matrix.
The runtime guarantees the matrix handed in here is the full
logical gradient: for pure-DP params the owner already holds
the full tensor; for TP-sharded params the TP gather
step (dmuon._backends.fsdp2.group.tp_gather_grads) has
reassembled the full matrix on the TP owner before this call.
direct_newton_schulz ¶
direct_newton_schulz(G: Tensor, steps: int = 5, eps: float = 1e-07, coefficients: Optional[list[list[float]]] = None) -> Tensor
Standard Newton-Schulz orthogonalization in direct (parameter) space.
Iterates on the full (m, n) matrix:
X_{k+1} = a_k X + b_k (X X^T) X + c_k (X X^T)^2 X
This is the classic formulation used by Muon/Moonlight. Compared to
:func:gram_newton_schulz (Gram-space), direct NS is simpler but:
- Does not benefit from SYRK symmetry acceleration
- Does not support the restart mechanism
- Intermediate ops are (m, n) instead of (m, m)
Use this when you want the standard algorithm without Gram-space optimizations, e.g., for baseline comparison or small matrices where SYRK overhead is not justified.
Args:
G: Gradient matrix (m, n), any dtype.
steps: Number of NS iterations (used only if coefficients is None).
eps: Normalization epsilon.
coefficients: Per-step (a, b, c) coefficients. Length determines
number of iterations. Defaults to :data:DEFAULT_COEFFICIENTS.
Returns: Orthogonalized update, same shape as G, in original dtype.
newton_schulz ¶
newton_schulz(G: Tensor, steps: int = 5, eps: float = 1e-07, coefficients: Optional[list[list[float]]] = None, restart_iterations: Optional[list[int]] = None) -> Tensor
Newton-Schulz orthogonalization (default: Gram-space backend).
Routes to :func:gram_newton_schulz by default for better precision
(per-step coefficients, restart mechanism, SYRK acceleration).
For the standard direct-space algorithm, use :func:direct_newton_schulz.
Args: G: Gradient matrix (m, n), any dtype. steps: Ignored (determined by len(coefficients)). eps: Normalization epsilon. coefficients: Per-step coefficients. Defaults to POLAR_EXPRESS_COEFFICIENTS. restart_iterations: Restart positions. Defaults to [2].
Returns: Orthogonalized update.
gram_newton_schulz ¶
gram_newton_schulz(G: Tensor, steps: int = 5, eps: float = 1e-07, coefficients: Optional[list[list[float]]] = None, restart_iterations: Optional[list[int]] = None, deterministic: bool = False) -> Tensor
Gram-space Newton-Schulz on a full (un-sharded) matrix.
Adapted from Dao-AILab/gram-newton-schulz. Iterates on the Gram matrix instead of the full gradient; uses per-step coefficients and restart mechanism for numerical stability.
TP handling: this function is TP-agnostic. For TP-sharded
parameters the runtime gathers the full matrix to a designated TP
owner via an All-to-All before calling this function (see
tp_design.md and dmuon._backends.fsdp2.group.tp_gather_grads).
There is no in-function TP all-reduce.
Args: G: Full gradient matrix (m, n), any dtype. steps: Ignored (determined by len(coefficients)). eps: Normalization epsilon. coefficients: Per-step coefficients. Defaults to POLAR_EXPRESS_COEFFICIENTS. restart_iterations: Iteration indices for restart. Defaults to [2]. deterministic: If True, use cuBLAS for all ops (no SYRK kernel).
Returns: Orthogonalized update, same shape as G.
gram_newton_schulz_factors ¶
gram_newton_schulz_factors(G: Tensor, steps: int = 5, eps: float = 1e-07, coefficients: Optional[list[list[float]]] = None, restart_iterations: Optional[list[int]] = None, deterministic: bool = False) -> tuple[tuple[Tensor, ...], bool, Tensor]
Return Gram-space NS left-factor segments for G.
gram_newton_schulz(G) applies one Gram multiplier segment per restart
window to the normalized and possibly transposed view of G. Returning
the segments instead of pre-composing them preserves the same fp16
matrix-multiplication order as the original implementation.
These factors are unrelated to attention Q projections. The internal
variable name Q follows the Gram-NS algebra, where the update is formed
by left-multiplying the normalized matrix by one or more Gram factors. The
TP path can broadcast these smaller Gram factors and let every TP rank
apply them to its local sequence/column shard instead of scattering the
full update.
Returns:
(factor_segments, transposed, normalizer) where normalizer is
G.norm()+eps under the same orientation rule as
:func:gram_newton_schulz.
gram_newton_schulz_q ¶
Backward-compatible alias for :func:gram_newton_schulz_factors.
The old name used Q from Gram-NS algebra. Prefer
gram_newton_schulz_factors in new code to avoid confusion with
attention Q projections.
gram_newton_schulz_distributed_local ¶
gram_newton_schulz_distributed_local(G_local: Tensor, group: Optional[ProcessGroup], *, transposed: bool, steps: int = 5, eps: float = 1e-07, coefficients: Optional[list[list[float]]] = None, restart_iterations: Optional[list[int]] = None) -> Tensor
Gram-space NS for a local column shard of the oriented matrix.
This is the math kernel for a Megatron-style TP shard-owner update path.
The full logical matrix is first oriented exactly like
:func:gram_newton_schulz: if transposed is true, G_local.T is the
local shard. The oriented local shard must be a column shard of the full
oriented matrix. The function reconstructs only the small Gram matrix via
TP all-reduce and returns the local shard of the full NS output.
It intentionally does not gather the full gradient, broadcast Q, or scatter a full update. The caller remains responsible for momentum, scaling, and weight decay.
gram_newton_schulz_distributed_local_batched ¶
gram_newton_schulz_distributed_local_batched(G_local: Tensor, group: Optional[ProcessGroup], *, transposed: bool, steps: int = 5, eps: float = 1e-07, coefficients: Optional[list[list[float]]] = None, restart_iterations: Optional[list[int]] = None) -> Tensor
Batched variant of :func:gram_newton_schulz_distributed_local.
G_local must be a 3D tensor [batch, rows, cols] containing local
TP shards with identical logical orientation. This helper is intentionally
narrow: it is the math primitive needed to replace many same-shaped
per-param Gram all-reduces with one packed all-reduce bucket.
gram_newton_schulz¶
TP-aware Gram NS with SYRK decomposition. See Tensor Parallelism.
gram_newton_schulz ¶
gram_newton_schulz(G: Tensor, steps: int = 5, eps: float = 1e-07, coefficients: Optional[list[list[float]]] = None, restart_iterations: Optional[list[int]] = None, deterministic: bool = False) -> Tensor
Gram-space Newton-Schulz on a full (un-sharded) matrix.
Adapted from Dao-AILab/gram-newton-schulz. Iterates on the Gram matrix instead of the full gradient; uses per-step coefficients and restart mechanism for numerical stability.
TP handling: this function is TP-agnostic. For TP-sharded
parameters the runtime gathers the full matrix to a designated TP
owner via an All-to-All before calling this function (see
tp_design.md and dmuon._backends.fsdp2.group.tp_gather_grads).
There is no in-function TP all-reduce.
Args: G: Full gradient matrix (m, n), any dtype. steps: Ignored (determined by len(coefficients)). eps: Normalization epsilon. coefficients: Per-step coefficients. Defaults to POLAR_EXPRESS_COEFFICIENTS. restart_iterations: Iteration indices for restart. Defaults to [2]. deterministic: If True, use cuBLAS for all ops (no SYRK kernel).
Returns: Orthogonalized update, same shape as G.
get_ns_backend¶
Inspect which NS kernel is active. Returns a one-line summary string
(e.g. "Gram NS · kernel=cute_sm80 (SM80, DMuon internal)", "Gram NS · kernel=quack (SM90, Tri Dao quack)", or "Gram NS · kernel=cublas (SM70, universal fallback)"). See
Backend dispatch.
get_ns_backend ¶
Return a human-readable one-liner describing the active NS kernel.
Format: "Gram NS · kernel=<name> (SM<ver>, <detail>)". Typical
results:
"Gram NS · kernel=cute_sm80 (SM80, DMuon internal)""Gram NS · kernel=quack (SM90, Tri Dao quack)""Gram NS · kernel=cublas (SM80, universal fallback)"
This is the terse one-liner meant for startup log / user sanity-check.
Use :func:get_backend_status for the full dict of availability
flags.
import dmuon print(dmuon.get_ns_backend()) Gram NS · kernel=cute_sm80 (SM80, DMuon internal)
get_backend_status¶
Full diagnostic dict of the NS kernel dispatch layer — sm_version,
auto_choice, and per-backend availability flags. Useful for
programmatic checks and bug reports.
get_backend_status ¶
Return a full diagnostic snapshot of the NS kernel dispatch state.
Returns a plain dict with:
sm_version— int, detected compute capability (0 on CPU)auto_choice— which backendkernel="auto"resolves toquack_available— bool, soft-dep flagcute_sm80_available— bool, CuteDSL SYRK importablecublas_always_available— bool, alwaysTrue
Useful for programmatic checks and bug reports.
YOU_COEFFICIENTS¶
5-step (a, b, c) coefficients from
@YouJiacheng.
YOU_COEFFICIENTS
module-attribute
¶
YOU_COEFFICIENTS = [[4.0848, -6.8946, 2.927], [3.9505, -6.3029, 2.6377], [3.7418, -5.5913, 2.3037], [2.8769, -3.1427, 1.2046], [2.8366, -3.0525, 1.2012]]
POLAR_EXPRESS_COEFFICIENTS¶
Default 5-step coefficients from Polar Express (arXiv:2505.16932, safety
factor 1.05). Used when no coefficients argument is provided.
POLAR_EXPRESS_COEFFICIENTS
module-attribute
¶
POLAR_EXPRESS_COEFFICIENTS = [(a / _SAFETY_FACTOR, b / _SAFETY_FACTOR ** 3, c / _SAFETY_FACTOR ** 5) for a, b, c in _UNMODIFIED_POLAR_EXPRESS]
Utilities — DMuon state management¶
no_sync¶
Context manager for gradient accumulation; suppresses DMuon reduce and FSDP2's reduce-scatter within the block. See Gradient Accumulation.
no_sync ¶
Context manager to disable gradient reduction for gradient accumulation.
Within this context, backward passes skip reduce communication and accumulate gradients locally. On the next backward outside this context, the accumulated gradients are merged and reduced normally.
This also disables FSDP2's gradient sync for symmetric parameters.
Usage::
for i, batch in enumerate(dataloader):
ctx = dmuon.no_sync(model) if (i + 1) % accum_steps != 0 else nullcontext()
with ctx:
loss = model(batch).loss / accum_steps
loss.backward()
if (i + 1) % accum_steps == 0:
optimizer.step()
optimizer.zero_grad()
prepare_muon_grads¶
Prepare all pending Muon gradients after backward. This is broader than a plain reduce wait because TP-sharded parameters may also need a TP gather before Muon can run.
prepare_muon_grads ¶
prepare_muon_grads(model: Module, *, use_reduce_stream: bool = False, coalesce_wait: bool = False) -> None
Prepare all pending Muon gradients after backward.
Gradient reduces are dispatched asynchronously during backward. This
function resolves those reduce tails and, for TP-sharded params, launches
the TP gather that materializes _tp_full_grad for the TP owner. It
drains all prepared groups before returning so callers can immediately
run Muon update code.
When every group shares one reduce stream, use_reduce_stream combined
with coalesce_wait lets advanced callers (such as the fused fast-clip
path) issue a single aggregate stream wait instead of draining each group
individually. Both default to False and preserve the per-group wait.
wait_all_reduces¶
Backward-compatible alias for prepare_muon_grads(). Called automatically by
Muon.step().
wait_all_reduces ¶
Backward-compatible alias for :func:prepare_muon_grads.
The historical name is kept for callers that explicitly wait after
loss.backward(). The operation is now broader than a reduce wait:
it also prepares TP gathered gradients needed by Muon.
broadcast_all_updates¶
Synchronous post-step replicate broadcast (HSDP Phase B). No-op in 1D mode. Prefer the async variant unless debugging.
broadcast_all_updates ¶
Sync post-step broadcast of updated _owned_data.
Dispatches on every dedicated group, then drains. FSDP2 path fans across the HSDP replicate axis (no-op for 1D); DDP path fans across the DP group. Use the same forward-order priority as the async path so sync/async modes enter collective-bearing groups in one deterministic sequence.
broadcast_all_updates_async¶
Async post-step replicate broadcast (default in Muon). Each layer's event
is consumed at the start of the next forward pass.
broadcast_all_updates_async ¶
Phase C.2 (async): dispatch the post-step replicate broadcasts on the dedicated replicate stream and return without waiting.
The wait is consumed per-group by the next forward iteration's
_pre_forward_wait hook (see DedicatedState._pre_forward),
letting each layer's replicate broadcast hide behind the compute of
the prior layers.
Group dispatch order follows the recorded forward order from the
previous iteration — Phase C.5 priority scheduling that mirrors
FSDP2's post_forward_order usage for backward prefetch
(_fsdp_param_group.py:469-474). First epoch falls through to
the model-walk order.
The async path keeps one pending event per group; the next forward drains the event before reading the group's data.
wait_all_replicate_broadcasts¶
Drain every group's pending async replicate broadcast. Call before reading
_owned_data outside the normal forward/step cycle.
wait_all_replicate_broadcasts ¶
Drain every group's pending async replicate broadcast.
Phase C.3 safety net for any code path that needs consistent
_owned_data without going through the forward hook — e.g.
get_model_state_dict / get_optimizer_state_dict in
:mod:dmuon.checkpoint call this before reading from global owners.
In sync / 1D mode this is a cheap no-op: every group is already IDLE.
wait_all_post_step_broadcasts¶
Compatibility alias for wait_all_replicate_broadcasts().
wait_all_post_step_broadcasts ¶
Alias for :func:wait_all_replicate_broadcasts.
Exposed under a more path-neutral name so DDP-path users do not have
to think about HSDP's replicate terminology. Both FSDP2-path
groups (HSDP replicate broadcast) and DDP-path groups (post-step
broadcast across the DP group) are drained.
clip_grad_norm_¶
Clip gradients for DMuon-owned Muon parameters.
clip_grad_norm_ ¶
clip_grad_norm_(optimizer: Any, max_norm: float | None, *, norm_type: float = 2.0, error_if_nonfinite: bool = False, foreach: bool | None = None, strategy: str | MuonGradClipStrategy = 'global_norm') -> MuonGradClipStats
Clip DMuon-owned Muon gradients only.
This does not touch AdamW/non-dedicated parameters. Training
frameworks should keep using PyTorch's native clipping for ordinary
param.grad tensors, and call this helper as the extra DMuon line.
register_muon_grad_clip_strategy¶
Register a custom strategy for clip_grad_norm_().
register_muon_grad_clip_strategy ¶
Register a custom Muon grad clip strategy.
The strategy receives live DMuon gradient tensors and may scale or inspect them in place. This is intentionally small so specialized schemes can be added without changing the public training-loop API.
MuonGradClipStats¶
Return type for DMuon gradient clipping.
MuonGradClipStats
dataclass
¶
MuonGradClipStats(total_norm: float, max_norm: float | None, norm_type: float, clip_coef: float, clipped: bool, param_count: int, found_inf: bool, strategy: str = 'global_norm')
Statistics returned by :func:clip_grad_norm_.
total_norm is the norm before clipping. max_norm=None means the
call was stats-only and did not scale any gradient.
clip_grad_norm_buckets_¶
Clip independent gradient buckets with segment-local semantics, using the optional CUDA fast path when available and pure Python otherwise. See Fast Gradient Clipping for the extension.
clip_grad_norm_buckets_ ¶
clip_grad_norm_buckets_(buckets: Sequence[GradClipBucket], max_norm: float | None, *, process_group=None, chunk_size: int | None = None) -> SegmentedGradClipResult
Clip gradient buckets with segment-local, distributed-safe semantics.
Each bucket keeps its own gradient norm and clip coefficient. The CUDA kernels accelerate the eligible tensors; ineligible tensors (or all of them when the extension is unavailable) use the torch reference math. The fast/slow choice is per tensor and purely local: it never changes the number or shape of collective operations, so this is safe to call from every rank in lockstep.
try_clip_optimizer_grad_norm_buckets_¶
Segmented regular / muon / adamw clip entry point for Wall-X-style
optimizers. Returns None when the CUDA extension is unavailable so callers
can fall back without any gradient mutation.
try_clip_optimizer_grad_norm_buckets_ ¶
try_clip_optimizer_grad_norm_buckets_(optimizer, *, regular_grads: Sequence[Tensor], adamw_grads: Sequence[Tensor], max_norm: float, process_group=None) -> SegmentedGradClipResult | None
Fast path for Wall-X style regular/Muon/dedicated-AdamW clipping.
Returns None only when the fast-clip CUDA extension is entirely
unavailable, letting the caller run its own path. Extension availability
is a property of the install and is therefore uniform across ranks, so this
None decision is rank-consistent and cannot desync collectives. Once a
result is produced, clipping goes through the distributed-safe
:func:clip_grad_norm_buckets_.
GradClipBucket¶
One independently clipped gradient segment.
GradClipBucket
dataclass
¶
One independently clipped gradient segment.
GradClipBucketStats¶
Per-segment pre-clip norm and coefficient.
GradClipBucketStats
dataclass
¶
GradClipBucketStats(name: str, total_norm: float, max_norm: float | None, clip_coef: float, clipped: bool, param_count: int, found_inf: bool, fastpath: bool)
Pre-clip norm and coefficient for one gradient segment.
SegmentedGradClipResult¶
Aggregate result of a multi-segment clip call.
SegmentedGradClipResult
dataclass
¶
SegmentedGradClipResult(total_norm: Tensor, bucket_stats: tuple[GradClipBucketStats, ...], fastpath: bool, fallback_reason: str | None = None)
Result for a multi-segment clipping call.
get_dedicated_params¶
Enumerate all DedicatedParam objects across the model.
get_dedicated_params ¶
Collect all DedicatedParam instances from a model.
get_owned_params¶
Filter DedicatedParam objects owned by a rank. Accepts int (1D) or
(shard, replicate) tuple (HSDP).
get_owned_params ¶
Collect DedicatedParam instances owned by a specific rank.
Accepts either a plain shard int (1D legacy form, matched against
owner_shard) or a (shard, replicate) tuple (matched against the
full owner_rank coord).
get_comm_ctx¶
get_comm_ctx ¶
get_comm_ctx(model: Module) -> Optional[DedicatedCommContext]
Get the DedicatedCommContext from a model, if it exists.
DedicatedCommContext¶
Shared CUDA streams (broadcast, reduce, replicate-broadcast) and prefetch
ordering state. Analogous to FSDP2's FSDPCommContext.
DedicatedCommContext ¶
DedicatedCommContext(device: device, replicate_group: Optional[ProcessGroup] = None, *, tp_buffer_reuse: bool | str = False, replicate_broadcast_bucket_mb: float = 0.0)
Shared communication context across all DedicatedParamGroups.
Streams:
broadcast_stream: high-priority stream for NCCL broadcast kernels on
the shard (dp_group) dimension.
reduce_stream: high-priority stream for NCCL reduce kernels on the
shard dimension.
replicate_reduce_stream: default-priority stream for HSDP Stage-2
replicate-axis reduce/all-reduce. It is intentionally separate
from the post-step publish stream so large AdamW-route all-reduces
do not sit in front of next-forward Muon publishes.
replicate_broadcast_stream: default-priority stream reserved for the
post-step inter-replicate-group broadcast used by HSDP-native Muon
and TP post-step scatter. Initialised here so Phase A code can
reference it without conditional guards; it stays idle until those
paths are enabled.
Process groups:
replicate_group: ProcessGroup spanning the replicate dimension of
the HSDP 2D mesh. None in 1D shard-only mode, which is the
Phase A default — downstream code falls back to the previous
single-dimension collectives when this is None.
Prefetch state: post_forward_order: records which groups ran forward, in order. Used in backward to determine the next group to prefetch.
Root post-backward fallback: all_states: every DedicatedState registers itself here so the autograd-engine callback can iterate and force-fire any group whose fast-path post-backward did not run (e.g., when no input tensor required gradient). post_backward_final_callback_queued: guards the callback so it is queued at most once per backward pass.
Rolling Stage-1 drain (1-outstanding):
last_reduced_group: most recently reduce_grads-dispatched group
whose Stage-1 shard reduce has not been safety-drained yet. Each
new _run_post_backward waits only on that Stage-1 event before
dispatching its own reduce, mirroring FSDP2's
reduce_scatter_state buffer-lifetime wait. HSDP Stage-2
replicate reduce/all-reduce is intentionally not waited here; it
is drained at the optimizer/root-post-backward boundary. Reset by
the root post-backward callback at the end of every backward pass.
Build the shared communication context for one model's dedicated-ownership groups.
Args:
device: CUDA device the streams and collectives run on.
replicate_group: ProcessGroup spanning the replicate
dimension of the HSDP 2D mesh. None in 1D shard-only
mode (the default); downstream code short-circuits the
replicate-dim reduce/broadcast when this is None.
tp_buffer_reuse: Whether TP gather/scatter should reuse per-param
scratch buffers. Accepts False/True or "gather",
"scatter", "all".
replicate_broadcast_bucket_mb: Optional HSDP post-step publish
bucket size in MiB. 0 keeps one coalesced publish per group.
Normally constructed by :func:dmuon.dedicate_params, not by
user code directly.
Diagnostics¶
The diagnostics helpers return rank-local, JSON-friendly summaries. They do not launch distributed collectives, so they are safe to call from benchmark logging code. Dump one summary per rank if you need a global view.
summarize_param_groups¶
Inspect how Muon routed trainable parameters across optimizer groups. Use
this after constructing the optimizer to check type-split routing, owner
counts, and whether param_policy selected muon, adamw, or
sharded_adamw plus the expected parameter/gradient dtypes.
import json
import dmuon
print(json.dumps(
dmuon.summarize_param_groups(model, optimizer),
indent=2,
default=str,
))
summarize_param_groups ¶
Summarize how dmuon.Muon routed trainable parameters.
The returned schema is intentionally JSON-friendly and stable enough for benchmark summaries. Counts are local-rank facts unless the key explicitly says it includes all dedicated params.
summarize_comm_plan¶
Inspect the DMuon communication plan for FSDP2/HSDP groups. The summary reports owner buckets, root ranks, route labels, and payload estimates. It is an estimate of planned tensor sizes, not a measured NCCL latency report.
summarize_comm_plan ¶
summarize_comm_plan(model: Module, *, max_groups: int = 64, max_params_per_group: int = 128) -> dict[str, Any]
Summarize DMuon communication roots and payload estimates.
This function mirrors the actual group ordering and owner buckets used by the FSDP2 backend. It is an estimate only: it reports planned tensor sizes and roots, not measured NCCL latency.
collect_forward_unshard_profile¶
Collect aggregate forward-unshard counters and CUDA-event timings from the
communication context. Enable collection before dedicate_params() creates
the context:
Then collect at a diagnostic boundary:
Use synchronize=True only at diagnostic boundaries. It calls
torch.cuda.synchronize(), so putting it inside a normal timing loop changes
the overlap behavior being measured.
collect_forward_unshard_profile ¶
collect_forward_unshard_profile(model: Module, *, clear: bool = True, synchronize: bool = False) -> dict[str, object]
Collect DMuon FSDP2 forward-unshard diagnostics for model.
The counters are only populated when DMUON_RECORD_FORWARD_PROFILE=1 is
set before :func:dmuon.dedicate_params builds the communication context.
synchronize=True should be used at diagnostic boundaries so CUDA-event
timings are complete; keep it disabled in normal measurement windows.
Checkpointing¶
All four are collective — every rank must call. Drains async state before reading/writing. Standard format: compatible with single-GPU and HuggingFace checkpoints. See Checkpointing.
get_model_state_dict¶
get_model_state_dict ¶
get_model_state_dict(model: Module, *, cpu_offload: bool = True, rank0_only: bool = True) -> dict[str, Tensor]
Get full model state dict with both dedicated and FSDP2 parameters.
Produces a state dict identical to what a single-GPU model would produce,
compatible with torch.save/torch.load and HuggingFace checkpoints.
For dedicated params: broadcasts _owned_data from each owner.
For FSDP2 symmetric params: all-gathers sharded DTensors.
Args:
model: Model with dedicate_params and fully_shard applied.
cpu_offload: Move tensors to CPU (default True, recommended for saving).
rank0_only: If True (default), only rank 0 returns a populated state
dict; other ranks return {}. All ranks still participate in
the NCCL broadcast / all_gather collectives (required for
correctness) but discard the gathered tensors so non-rank0 nodes
don't accumulate ~tens of GB of CPU RAM per rank. Matches the
semantics of FSDP1's FullStateDictConfig(rank0_only=True).
Returns:
Complete state dict with full (unsharded) tensors for all parameters
on rank 0; empty dict on other ranks when rank0_only is True.
set_model_state_dict¶
set_model_state_dict ¶
Load a full state dict into a DMuon model.
Handles both dedicated params (copy to owner's _owned_data) and
FSDP2 symmetric params (manual sharding into DTensors).
The state dict should contain full (unsharded) tensors, as produced by
:func:get_model_state_dict or a single-GPU model.state_dict().
Args:
model: Model with dedicate_params and fully_shard applied.
state_dict: Full state dict mapping FQN to unsharded tensors.
get_optimizer_state_dict¶
get_optimizer_state_dict ¶
get_optimizer_state_dict(model: Module, optimizer: Any, *, cpu_offload: bool = True, rank0_only: bool = True) -> dict
Get optimizer state dict for a DMuon Muon optimizer.
Produces a dict with four state sections:
- "fsdp": FSDP2 AdamW state (FQN-keyed, full tensors)
- "dedicated": Muon momentum buffers (FQN-keyed, broadcast from owners)
- "dedicated_adamw": dedicated AdamW state for owner-managed base params
- "param_groups": Hyperparameters for both groups
Args:
model: Model with dedicate_params and fully_shard applied.
optimizer: :class:dmuon.Muon optimizer instance.
cpu_offload: Move tensors to CPU (default True).
rank0_only: If True (default), only rank 0 returns a populated dict;
other ranks return an empty {"fsdp": {}, "dedicated": {},
"param_groups": []}. All ranks still participate in the NCCL
broadcast / all_gather collectives.
Returns: Optimizer state dict in DMuon format.
set_optimizer_state_dict¶
set_optimizer_state_dict ¶
Load optimizer state dict into a DMuon Muon optimizer.
Args:
model: Model with dedicate_params and fully_shard applied.
optimizer: :class:dmuon.Muon optimizer instance.
state_dict: Optimizer state dict as produced by
:func:get_optimizer_state_dict.