API 文档¶

TL;DR

DMuon 公开五个功能区：初始化（dedicate_params、install_patch）、 优化器（Muon、NewtonSchulz、NS 函数与常量）、状态管理 （no_sync、wait_all_reduces、replicate-broadcast 工具函数、 DedicatedCommContext）、诊断（summarize_param_groups、 summarize_comm_plan、collect_forward_unshard_profile）以及检查点 （get/set_model/optimizer_state_dict）。从 dedicate_params + Muon 开始；需要精细控制时再使用其余接口。

初始化¶

dedicate_params¶

在 fully_shard() 之前调用一次。将每个 Muon 目标参数分配给单一 owner rank，并注册逐层的前向/反向 hook。常见自定义点见自定义 Hook 边界和 Z2 与 Z3 模式。

默认情况下，dedicate_params() 使用 process_group_policy="isolated"。 DMuon 会为自己的 DP/HSDP/TP collective 复制一套 mesh process group，避免外部 trainer 的 logging、metrics、checkpoint collective 插入 DMuon 异步 post-step publish 所在的 NCCL communicator 序列。只有在明确需要复用调用方 DeviceMesh group 的历史行为时，才传 process_group_policy="shared"。

isolated 只改变 process group 的归属关系，默认不会在 step 末尾加 barrier，因此 DMuon 的异步 publish 仍然可以和下一步训练 overlap。只有在排查疑似 process group 顺序问题时，才设置 DMUON_ISOLATED_PG_BARRIER=1，强制对 DMuon-owned groups 做严格的 post-step fence。性能测试默认不要打开这个 fence，除非实验目标就是测 no-overlap 行为。

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 路径下的 dedicated parameter 初始化入口，用于 data-parallel 模型保持 replicated、而不是由 FSDP2 分片的场景。

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.Parameter stays live.
Does not register forward-broadcast / reshard hooks. DedicatedState still attaches hooks at the layer boundary, but unshard / reshard degrade to no-op on DDP groups; the real work is reduce_grads after backward + post_step_broadcast after optim.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 路径下启用 TP 的 dedicated parameter 初始化入口，适用于每个 replicated data-parallel group 内部还有 Tensor Parallelism 的场景。

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。

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 ¶

replicate(model: Module, mesh: DeviceMesh) -> Module

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 ¶

replicate_tp(model: Module, mesh: DeviceMesh) -> Module

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 ¶

get_replicated_group(model: Module) -> Optional[_ReplicatedGroup]

Return the _ReplicatedGroup attached to model, or None.

replicate_tp¶

dedicate_params_ddp_tp() 的 TP-aware companion，用于非专属 DTensor 参数。

replicate_tp ¶

replicate_tp(model: Module, mesh: DeviceMesh) -> Module

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 会自动调用此函数。除非在不经过正常 import 路径的情况下构建 DMuon 环境，否则无需手动调用。

install_patch ¶

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.

优化器¶

Muon¶

主优化器类。在同一对象中管理 matrix-routed 专属参数上的 Muon （Newton-Schulz + 动量）和 base path 上的 AdamW。Base path 可以是普通 FSDP2-managed 参数，也可以是通过 param_policy 或旧版 route_hint_fn 选择的 DMuon-managed sharded AdamW 参数。兼容 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 ¶

consume_last_step_profile() -> dict[str, object]

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 ¶

step(closure=None)

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 ¶

zero_grad(set_to_none: bool = True)

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¶

可配置的 NS 后端对象。传入 Muon(ns_backend=...) 以选择算法变体、覆盖系数或启用确定性模式。完整对比见 Newton-Schulz 变体。

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 ¶

local(G: Tensor, steps: int) -> Tensor

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¶

独立 NS 函数，默认路由至 Gram 空间后端。在优化器循环外需要 NS 时直接使用，例如自定义训练循环或实验代码。

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 ¶

local(G: Tensor, steps: int) -> Tensor

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 ¶

gram_newton_schulz_q(*args, **kwargs)

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 感知的 Gram NS（带 SYRK 分解）。Muon 内部为 Tensor-Parallel 参数调用此函数；此处暴露供构建自定义 TP 优化器的用户使用。见张量并行。

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¶

查询当前活跃的 NS 内核。返回形如 "Gram NS · kernel=cute_sm80 (SM80, DMuon internal)"、 "Gram NS · kernel=quack (SM90, Tri Dao quack)" 或 "Gram NS · kernel=cublas (SM70, universal fallback)" 的单行摘要。详见后端分发。

get_ns_backend ¶

get_ns_backend() -> str

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¶

NS 内核分发层的完整诊断 dict —— sm_version、auto_choice，以及各后端的可用性标志。适合程序化检查和 bug report。

get_backend_status ¶

get_backend_status() -> dict

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 backend kernel="auto" resolves to
quack_available — bool, soft-dep flag
cute_sm80_available — bool, CuteDSL SYRK importable
cublas_always_available — bool, always True

Useful for programmatic checks and bug reports.

YOU_COEFFICIENTS¶

来自 @YouJiacheng 的 5 步逐迭代 (a, b, c) 系数。传入 NewtonSchulz(coefficients=...) 或直接传入 NS 函数的 coefficients 参数。

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¶

默认 5 步系数，来自 Polar Express 论文（arXiv:2505.16932），应用了 1.05 安全因子。不传 coefficients 参数时默认使用。

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]

工具函数 — DMuon 状态管理¶

no_sync¶

梯度累积的上下文管理器。在上下文内抑制 DMuon reduce 和 FSDP2 的 reduce-scatter；最后一个 micro-batch 在上下文外调用 backward 以触发合并 reduce。见梯度累积。

no_sync ¶

no_sync(model: Module)

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¶

在 backward 之后准备所有 pending 的 Muon 梯度。它不只是等待 reduce；对 TP-sharded 参数，还可能需要在 Muon 运行前触发 TP gather。

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¶

prepare_muon_grads() 的向后兼容 alias。Muon.step() 会自动调用；仅当需要在 backward 和 step 之间手动访问 prepared gradient 时才需要单独调用。

wait_all_reduces ¶

wait_all_reduces(model: Module) -> None

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¶

同步的后置 replicate broadcast（HSDP Phase B 路径）。将更新后的 _owned_data 从全局 owner 广播到每个 replicate 对等节点。在 1D shard-only 模式下为空操作。除非调试，优先使用异步变体。

broadcast_all_updates ¶

broadcast_all_updates(model: Module) -> None

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¶

异步的后置 replicate broadcast（Muon 默认值）。立即返回；每一层的 event 在下一次前向传播开始时被消费。

broadcast_all_updates_async ¶

broadcast_all_updates_async(model: Module) -> None

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¶

等待所有 group 的异步 replicate broadcast 完成。在正常前向/step 周期之外需要读取 _owned_data 的代码（如自定义检查点或评估逻辑）中调用。

wait_all_replicate_broadcasts ¶

wait_all_replicate_broadcasts(model: Module) -> None

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¶

wait_all_replicate_broadcasts() 的兼容 alias。

wait_all_post_step_broadcasts ¶

wait_all_post_step_broadcasts(model: Module) -> None

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_¶

裁剪 DMuon-owned Muon 参数的梯度。

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¶

为 clip_grad_norm_() 注册自定义策略。

register_muon_grad_clip_strategy ¶

register_muon_grad_clip_strategy(name: str, strategy: MuonGradClipStrategy) -> None

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¶

DMuon 梯度裁剪的返回类型。

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_¶

以分段局部语义裁剪相互独立的梯度分段；可用时走可选的 CUDA 快速路径，否则回退纯 Python。扩展的编译见快速梯度裁剪。

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_¶

面向 Wall-X 式优化器的 regular / muon / adamw 分段裁剪入口。CUDA 扩展不可用时返回 None，调用方可在不触碰梯度的情况下回退。

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¶

一个可独立裁剪的梯度分段。

GradClipBucket `dataclass` ¶

GradClipBucket(name: str, grads: Sequence[Tensor], reduce: bool = True)

One independently clipped gradient segment.

GradClipBucketStats¶

单个分段裁剪前的范数与系数。

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¶

多分段裁剪调用的聚合结果。

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¶

枚举模型中所有 DedicatedParam 对象。用于检查 ownership 分配、参数数量和负载均衡情况。

get_dedicated_params ¶

get_dedicated_params(model: Module) -> list[DedicatedParam]

Collect all DedicatedParam instances from a model.

get_owned_params¶

筛选属于指定 rank 坐标的 DedicatedParam 对象。接受整数（1D）或 (shard, replicate) 元组（HSDP）。

get_owned_params ¶

get_owned_params(model: Module, rank: OwnerRankLike) -> list[DedicatedParam]

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¶

获取存储在模型上的 DedicatedCommContext。若未调用 dedicate_params 则返回 None。

get_comm_ctx ¶

get_comm_ctx(model: Module) -> Optional[DedicatedCommContext]

Get the DedicatedCommContext from a model, if it exists.

DedicatedCommContext¶

持有专属 CUDA stream（broadcast、reduce、replicate-broadcast）和预取顺序状态的共享通信上下文。类比 FSDP2 的 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.

reset_post_forward_order ¶

reset_post_forward_order() -> None

Clear post-forward order. Call at the start of each forward pass.

consume_forward_unshard_profile ¶

consume_forward_unshard_profile(*, clear: bool = True) -> dict[str, object]

Return aggregate CUDA-event timings for forward unshard diagnostics.

诊断¶

诊断函数返回当前 rank 的 JSON-friendly 摘要。它们不会发起 distributed collective，因此可以放在 benchmark 日志代码里。若需要全局视图，让每个 rank 各自 dump，再由外部脚本汇总。

summarize_param_groups¶

检查 Muon 如何把可训练参数路由到 optimizer groups。构造 optimizer 后调用它，可以确认 type-split routing、owner 数量，以及 param_policy 是否按预期选择了 muon、adamw 或 sharded_adamw，同时确认参数和梯度 dtype。

import json
import dmuon

print(json.dumps(
    dmuon.summarize_param_groups(model, optimizer),
    indent=2,
    default=str,
))

summarize_param_groups ¶

summarize_param_groups(model: Module, optimizer: Any, *, max_rows: int = 128) -> dict[str, Any]

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¶

检查 FSDP2/HSDP groups 的 DMuon 通信计划。摘要会报告 owner buckets、root ranks、route 标签和 payload 估算。它表示计划中的 tensor 大小，不是实测 NCCL latency。

print(json.dumps(
    dmuon.summarize_comm_plan(model),
    indent=2,
    default=str,
))

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¶

从通信上下文里收集 forward-unshard 计数器和 CUDA-event timing。必须在 dedicate_params() 创建通信上下文前打开采集：

DMUON_RECORD_FORWARD_PROFILE=1 torchrun ...

然后在诊断边界读取：

profile = dmuon.collect_forward_unshard_profile(
    model,
    synchronize=True,
)

synchronize=True 只适合诊断边界。它会调用 torch.cuda.synchronize()；如果放进正常 step timing loop，会改变正在测量的 overlap 行为。

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.

检查点¶

以下四个函数均为集体操作 — 每个 rank 都必须调用。它们在读写张量前会排空待处理的异步状态。状态字典为标准格式，兼容单 GPU 的 torch.save/torch.load 和 HuggingFace 检查点。见检查点指南。

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 ¶

set_model_state_dict(model: Module, state_dict: dict[str, Tensor]) -> None

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 ¶

set_optimizer_state_dict(model: Module, optimizer: Any, state_dict: dict) -> None

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.

API 文档¶

初始化¶

dedicate_params¶

dedicate_params ¶

dedicate_params_ddp¶

dedicate_params_ddp ¶

dedicate_params_ddp_tp¶

dedicate_params_ddp_tp ¶

replicate¶

replicate ¶

replicate ¶

replicate_tp ¶

get_replicated_group ¶

replicate_tp¶

replicate_tp ¶

install_patch¶

install_patch ¶

优化器¶

Muon¶

Muon ¶

last_muon_grad_clip_stats property ¶

consume_last_step_profile ¶

clip_grad_norm_ ¶

step ¶

zero_grad ¶

NewtonSchulz¶

NewtonSchulz ¶

local ¶

newton_schulz¶

newton_schulz ¶

NewtonSchulz ¶

local ¶

direct_newton_schulz ¶

newton_schulz ¶

gram_newton_schulz ¶

gram_newton_schulz_factors ¶

gram_newton_schulz_q ¶

gram_newton_schulz_distributed_local ¶

gram_newton_schulz_distributed_local_batched ¶

gram_newton_schulz¶

gram_newton_schulz ¶

get_ns_backend¶

get_ns_backend ¶

get_backend_status¶

get_backend_status ¶

YOU_COEFFICIENTS¶

YOU_COEFFICIENTS module-attribute ¶

POLAR_EXPRESS_COEFFICIENTS¶

POLAR_EXPRESS_COEFFICIENTS module-attribute ¶

工具函数 — DMuon 状态管理¶

no_sync¶

no_sync ¶

prepare_muon_grads¶

prepare_muon_grads ¶

wait_all_reduces¶

wait_all_reduces ¶

broadcast_all_updates¶

broadcast_all_updates ¶

broadcast_all_updates_async¶

broadcast_all_updates_async ¶

wait_all_replicate_broadcasts¶

wait_all_replicate_broadcasts ¶

wait_all_post_step_broadcasts¶

wait_all_post_step_broadcasts ¶

clip_grad_norm_¶

clip_grad_norm_ ¶

register_muon_grad_clip_strategy¶

register_muon_grad_clip_strategy ¶

MuonGradClipStats¶

MuonGradClipStats dataclass ¶

clip_grad_norm_buckets_¶

clip_grad_norm_buckets_ ¶

try_clip_optimizer_grad_norm_buckets_¶

try_clip_optimizer_grad_norm_buckets_ ¶

GradClipBucket¶

GradClipBucket dataclass ¶

GradClipBucketStats¶

GradClipBucketStats dataclass ¶

SegmentedGradClipResult¶

SegmentedGradClipResult dataclass ¶

last_muon_grad_clip_stats `property` ¶

YOU_COEFFICIENTS `module-attribute` ¶

POLAR_EXPRESS_COEFFICIENTS `module-attribute` ¶

MuonGradClipStats `dataclass` ¶

GradClipBucket `dataclass` ¶

GradClipBucketStats `dataclass` ¶

SegmentedGradClipResult `dataclass` ¶