diff --git a/examples/v1/config/sft_qwen3p5.py b/examples/v1/config/sft_qwen3p5.py index 9075f4953d..637c28b469 100644 --- a/examples/v1/config/sft_qwen3p5.py +++ b/examples/v1/config/sft_qwen3p5.py @@ -102,6 +102,7 @@ def _get_bool_env(name: str, default: bool = False) -> bool: optim_cfg = MuonConfig( lr=float(os.environ.get("LR", "2e-5")), weight_decay=float(os.environ.get("WEIGHT_DECAY", "0.05")), + use_gram_newton_schulz=_get_bool_env("USE_GRAM_NEWTON_SCHULZ", False), ) lr_cfg = LRConfig( lr_type="cosine", diff --git a/pyproject.toml b/pyproject.toml index a941493dfe..a5fc32a482 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -81,6 +81,11 @@ video = [ "decord", "av", ] +gram-ns = [ + "torch>=2.7.1", + "quack-kernels==0.5.0", + "nvidia-cutlass-dsl==4.5.2", +] all = [ "jsonlines", "decord", diff --git a/tests/optim/test_gram_newton_schulz.py b/tests/optim/test_gram_newton_schulz.py new file mode 100644 index 0000000000..c9cfa192fb --- /dev/null +++ b/tests/optim/test_gram_newton_schulz.py @@ -0,0 +1,107 @@ +# Copyright (c) OpenMMLab. All rights reserved. + +from typing import Any + +import pytest +import torch + +from xtuner.v1.config.optim import MuonConfig +from xtuner.v1.optim import Muon +from xtuner.v1.optim import gram_newton_schulz as gram_ns_module +from xtuner.v1.optim.gram_newton_schulz import ( + POLAR_EXPRESS_COEFFICIENTS, + POLAR_EXPRESS_SAFETY_FACTOR, + GramNewtonSchulz, +) + + +def _cpu_orthogonalizer(**kwargs: Any) -> GramNewtonSchulz: + return GramNewtonSchulz(use_kernels=False, compile_kwargs=None, **kwargs) + + +class TestGramNewtonSchulz: + def test_polar_express_coefficients_include_recommended_safety_factor(self) -> None: + unmodified_first = (8.28721201814563, -23.595886519098837, 17.300387312530933) + expected = ( + unmodified_first[0] / POLAR_EXPRESS_SAFETY_FACTOR, + unmodified_first[1] / POLAR_EXPRESS_SAFETY_FACTOR**3, + unmodified_first[2] / POLAR_EXPRESS_SAFETY_FACTOR**5, + ) + + assert len(POLAR_EXPRESS_COEFFICIENTS) == 5 + torch.testing.assert_close(torch.tensor(POLAR_EXPRESS_COEFFICIENTS[0]), torch.tensor(expected)) + + def test_adaptive_selection_falls_back_for_square_and_low_aspect_ratio(self) -> None: + orthogonalizer = _cpu_orthogonalizer(min_aspect_ratio=2.0) + + assert orthogonalizer._selection(256, 256) == ("standard_fallback", "torch") + assert orthogonalizer._selection(256, 512) == ("standard_fallback", "torch") + assert orthogonalizer._selection(256, 768) == ("gram", "torch") + assert orthogonalizer._selection(768, 256) == ("gram", "torch") + + orthogonalizer._kernel_backend = object() + assert orthogonalizer._selection(256, 768) == ("gram", "cutedsl") + + def test_gram_matches_recommended_standard_path_on_cpu(self) -> None: + torch.manual_seed(0) + inputs = (torch.randn(12, 4), torch.randn(4, 12)) + gram = _cpu_orthogonalizer(min_aspect_ratio=2.0) + standard = _cpu_orthogonalizer(min_aspect_ratio=100.0) + + for input_ in inputs: + gram_output = gram(input_) + standard_output = standard(input_) + assert gram_output.dtype == input_.dtype + assert torch.isfinite(gram_output).all() + torch.testing.assert_close(gram_output, standard_output, atol=1e-2, rtol=1e-2) + + def test_concatenated_experts_are_orthogonalized_independently(self) -> None: + torch.manual_seed(1) + orthogonalizer = _cpu_orthogonalizer(min_aspect_ratio=2.0) + experts = torch.randn(3, 12, 4, dtype=torch.float32) + concatenated = experts.view(-1, experts.size(-1)) + + batched = orthogonalizer(concatenated, num_experts=experts.size(0)) + independent = torch.cat([orthogonalizer(expert) for expert in experts]) + + assert batched.shape == concatenated.shape + assert batched.dtype == concatenated.dtype + torch.testing.assert_close(batched, independent) + + def test_rejects_invalid_expert_layout(self) -> None: + orthogonalizer = _cpu_orthogonalizer() + + with pytest.raises(ValueError, match="Invalid num_experts"): + orthogonalizer(torch.randn(12, 4), num_experts=5) + + +class TestMuonGramNewtonSchulz: + def test_config_keeps_gram_ns_opt_in(self) -> None: + config = MuonConfig() + + assert config.use_gram_newton_schulz is False + assert config.gram_restart_iterations == (2,) + assert config.gram_ns_epsilon == 1e-7 + assert config.gram_ns_min_aspect_ratio == 2.0 + assert config.gram_ns_torch_compile is True + + def test_switch_selects_gram_ns_callback_without_cuda(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(gram_ns_module, "_make_kernel_backend", lambda: None) + parameter = torch.nn.Parameter(torch.randn(12, 4)) + + optimizer = Muon( + [parameter], + use_gram_newton_schulz=True, + gram_restart_iterations=(2,), + gram_ns_min_aspect_ratio=2.0, + gram_ns_torch_compile=False, + ) + + assert isinstance(optimizer._newton_schulz_func, GramNewtonSchulz) + assert optimizer._newton_schulz_func.restart_iterations == (2,) + + def test_switch_rejects_other_newton_schulz_implementations(self) -> None: + parameter = torch.nn.Parameter(torch.randn(12, 4)) + + with pytest.raises(ValueError, match="cannot be combined"): + Muon([parameter], use_gram_newton_schulz=True, use_triton=True) diff --git a/xtuner/v1/config/optim.py b/xtuner/v1/config/optim.py index 4ad9881b3c..73557420f1 100644 --- a/xtuner/v1/config/optim.py +++ b/xtuner/v1/config/optim.py @@ -78,6 +78,21 @@ class MuonConfig(OptimConfig): enable_all2all: Annotated[ bool, Parameter(help="Allow all-to-all comm strategy; set False on topologies where all-to-all is unreliable") ] = True + use_gram_newton_schulz: Annotated[ + bool, Parameter(help="Enable CuTeDSL Gram Newton-Schulz with automatic standard-NS fallback") + ] = False + gram_ns_epsilon: Annotated[ + float, Parameter(help="Normalization epsilon used by Gram NS and its standard fallback") + ] = 1e-7 + gram_restart_iterations: Annotated[ + Tuple[int, ...], Parameter(help="Zero-based Gram NS iterations before which to restart") + ] = (2,) + gram_ns_min_aspect_ratio: Annotated[ + float, Parameter(help="Use Gram NS only when a matrix aspect ratio is greater than this value") + ] = 2.0 + gram_ns_torch_compile: Annotated[ + bool, Parameter(help="Compile Gram NS with torch.compile fullgraph/reduce-overhead mode") + ] = True def build(self, model): trainable_parameters_names = model.trainable_parameters() @@ -205,6 +220,11 @@ def build(self, model): flatten=True, adjust_lr=self.adjust_lr, use_triton=False, + use_gram_newton_schulz=self.use_gram_newton_schulz, + gram_ns_epsilon=self.gram_ns_epsilon, + gram_restart_iterations=self.gram_restart_iterations, + gram_ns_min_aspect_ratio=self.gram_ns_min_aspect_ratio, + gram_ns_torch_compile=self.gram_ns_torch_compile, epsilon=self.eps, enable_all2all=self.enable_all2all, muon_split_sizes=muon_split_sizes, diff --git a/xtuner/v1/optim/gram_newton_schulz.py b/xtuner/v1/optim/gram_newton_schulz.py new file mode 100644 index 0000000000..a4a1e249e6 --- /dev/null +++ b/xtuner/v1/optim/gram_newton_schulz.py @@ -0,0 +1,249 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""Gram Newton-Schulz orthogonalization for XTuner's distributed Muon. + +The algorithm and CuTeDSL-backed GEMM interface are adapted from +https://github.com/Dao-AILab/gram-newton-schulz by Jack Zhang, Noah Amsel, +Berlin Chen, and Tri Dao (MIT license). +""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from types import SimpleNamespace +from typing import Any + +import torch +from torch import Tensor + + +SYMMETRIC_KERNEL_TILE_SIZE = 256 +POLAR_EXPRESS_SAFETY_FACTOR = 1.05 +_UNMODIFIED_POLAR_EXPRESS_COEFFICIENTS = ( + (8.28721201814563, -23.595886519098837, 17.300387312530933), + (4.107059111542203, -2.9478499167379106, 0.5448431082926601), + (3.9486908534822946, -2.908902115962949, 0.5518191394370137), + (3.3184196573706015, -2.488488024314874, 0.51004894012372), + (2.300652019954817, -1.6689039845747493, 0.4188073119525673), +) +POLAR_EXPRESS_COEFFICIENTS = tuple( + ( + a / POLAR_EXPRESS_SAFETY_FACTOR, + b / POLAR_EXPRESS_SAFETY_FACTOR**3, + c / POLAR_EXPRESS_SAFETY_FACTOR**5, + ) + for a, b, c in _UNMODIFIED_POLAR_EXPRESS_COEFFICIENTS +) + +_TORCH_BACKEND = SimpleNamespace( + sym_mm=lambda A, B: A @ B, + sym_baddbmm=lambda A, B, C, alpha=1.0, beta=1.0: torch.baddbmm(C, A, B, alpha=alpha, beta=beta), + mm=lambda A, B: A @ B, + mm_add=lambda A, B, C, beta: torch.baddbmm(C, A, B, beta=beta), +) + + +def _make_kernel_backend() -> SimpleNamespace: + try: + from quack.gemm_interface import gemm, gemm_add, gemm_symmetric + except ImportError as error: + raise ImportError( + "CuTeDSL Gram Newton-Schulz requires quack-kernels and nvidia-cutlass-dsl. " + "Install XTuner with `pip install -e '.[gram-ns]' --no-build-isolation`." + ) from error + + return SimpleNamespace( + sym_mm=gemm_symmetric, + sym_baddbmm=lambda A, B, C, alpha=1.0, beta=1.0: gemm_symmetric(A, B, C=C, alpha=alpha, beta=beta), + mm=lambda A, B: gemm(A, B, tuned=False), + mm_add=lambda A, B, C, beta: gemm_add(A, B, C=C, beta=beta, tuned=False), + ) + + +def _make_gram_newton_schulz( + ops: SimpleNamespace, + coefficients: Sequence[tuple[float, float, float]], + restart_iterations: Sequence[int], + epsilon: float, + compile_kwargs: dict[str, Any] | None, +) -> Callable[[Tensor], Tensor]: + """Build a closure with fixed operators so ``torch.compile`` sees a stable + graph.""" + coefficients = tuple(coefficients) + restart_set = frozenset(restart_iterations) + + def gram_newton_schulz(X: Tensor) -> Tensor: + tall = X.size(-2) > X.size(-1) + X = X.to(torch.float32) + X = X / (X.norm(dim=(-2, -1), keepdim=True) + epsilon) + X = X.to(torch.float16) + + R = ops.sym_mm(X.mT, X) if tall else ops.sym_mm(X, X.mT) + identity = ( + torch.eye(R.size(-1), device=X.device, dtype=X.dtype).unsqueeze(0).expand(R.size(0), -1, -1).contiguous() + ) + Q = None + + for iteration, (a, b, c) in enumerate(coefficients): + if iteration in restart_set and iteration != 0: + X = ops.mm(X, Q) if tall else ops.mm(Q, X) + R = ops.sym_mm(X.mT, X) if tall else ops.sym_mm(X, X.mT) + Q = None + + Z = ops.sym_baddbmm(R, R, C=R, alpha=c, beta=b) + if iteration == 0 or iteration in restart_set: + Q = Z + a * identity + else: + Q = ops.sym_baddbmm(Q, Z, C=Q, beta=a) + + if iteration < len(coefficients) - 1 and iteration + 1 not in restart_set: + RZ = ops.sym_baddbmm(R, Z, C=R, beta=a) + R = ops.sym_baddbmm(Z, RZ, C=RZ, beta=a) + + return ops.mm(X, Q) if tall else ops.mm(Q, X) + + if compile_kwargs is not None: + gram_newton_schulz = torch.compile(gram_newton_schulz, **compile_kwargs) + return gram_newton_schulz + + +def _make_standard_newton_schulz( + ops: SimpleNamespace, + coefficients: Sequence[tuple[float, float, float]], + epsilon: float, + compile_kwargs: dict[str, Any] | None, +) -> Callable[[Tensor], Tensor]: + """Build the recommended standard NS fallback with the same FP16 policy.""" + coefficients = tuple(coefficients) + + def standard_newton_schulz(X: Tensor) -> Tensor: + tall = X.size(-2) > X.size(-1) + X = X.to(torch.float32) + X = X / (X.norm(dim=(-2, -1), keepdim=True) + epsilon) + X = X.to(torch.float16) + + for a, b, c in coefficients: + A = ops.sym_mm(X.mT, X) if tall else ops.sym_mm(X, X.mT) + B = ops.sym_baddbmm(A, A, C=A, alpha=c, beta=b) + X = ops.mm_add(X, B, C=X, beta=a) if tall else ops.mm_add(B, X, C=X, beta=a) + return X + + if compile_kwargs is not None: + standard_newton_schulz = torch.compile(standard_newton_schulz, **compile_kwargs) + return standard_newton_schulz + + +class GramNewtonSchulz: + """XTuner-compatible adaptive Gram Newton-Schulz callable. + + Rectangular matrices whose aspect ratio is greater than + ``min_aspect_ratio`` use Gram NS. Square and low-aspect-ratio matrices use + standard NS. Both paths use FP32 normalization followed by FP16 internal + computation and the recommended safety-scaled Polar Express coefficients. + + Args: + epsilon (float): Normalization epsilon used by both algorithms. + use_kernels (bool): Whether to use CuTeDSL kernels for supported matrix shapes. + restart_iterations (tuple[int, ...]): Iterations before which Gram NS rebuilds the Gram matrix. + min_aspect_ratio (float): Gram NS is used only above this matrix aspect ratio. + compile_kwargs (dict[str, Any] | None): Keyword arguments passed to ``torch.compile``. ``None`` disables + compilation. + """ + + def __init__( + self, + *, + epsilon: float = 1e-7, + use_kernels: bool = True, + restart_iterations: tuple[int, ...] = (2,), + min_aspect_ratio: float = 2.0, + compile_kwargs: dict[str, Any] | None = None, + ) -> None: + if epsilon <= 0: + raise ValueError(f"epsilon must be positive, got {epsilon}") + if min_aspect_ratio < 1.0: + raise ValueError(f"min_aspect_ratio must be at least 1, got {min_aspect_ratio}") + if any(iteration <= 0 or iteration >= len(POLAR_EXPRESS_COEFFICIENTS) for iteration in restart_iterations): + raise ValueError( + f"restart_iterations must be in [1, {len(POLAR_EXPRESS_COEFFICIENTS) - 1}], got {restart_iterations}" + ) + + self.epsilon = epsilon + self.restart_iterations = tuple(restart_iterations) + self.min_aspect_ratio = min_aspect_ratio + self._compiled = compile_kwargs is not None + + kernel_backend = _make_kernel_backend() if use_kernels else None + self._kernel_backend = kernel_backend + self._gram_kernel: Callable[[Tensor], Tensor] | None = None + self._standard_kernel: Callable[[Tensor], Tensor] | None = None + self._gram_torch = _make_gram_newton_schulz( + _TORCH_BACKEND, + POLAR_EXPRESS_COEFFICIENTS, + self.restart_iterations, + epsilon, + compile_kwargs, + ) + self._standard_torch = _make_standard_newton_schulz( + _TORCH_BACKEND, POLAR_EXPRESS_COEFFICIENTS, epsilon, compile_kwargs + ) + if kernel_backend is not None: + self._gram_kernel = _make_gram_newton_schulz( + kernel_backend, + POLAR_EXPRESS_COEFFICIENTS, + self.restart_iterations, + epsilon, + compile_kwargs, + ) + self._standard_kernel = _make_standard_newton_schulz( + kernel_backend, POLAR_EXPRESS_COEFFICIENTS, epsilon, compile_kwargs + ) + + def _selection(self, rows: int, cols: int) -> tuple[str, str]: + aspect_ratio = max(rows, cols) / min(rows, cols) + algorithm = "gram" if rows != cols and aspect_ratio > self.min_aspect_ratio else "standard_fallback" + backend = ( + "cutedsl" + if self._kernel_backend is not None and min(rows, cols) >= SYMMETRIC_KERNEL_TILE_SIZE + else "torch" + ) + return algorithm, backend + + def _run_selected(self, X: Tensor) -> Tensor: + rows, cols = X.shape[-2:] + algorithm, backend = self._selection(rows, cols) + if backend == "cutedsl": + kernel = self._gram_kernel if algorithm == "gram" else self._standard_kernel + assert kernel is not None + output = kernel(X) + else: + torch_implementation = self._gram_torch if algorithm == "gram" else self._standard_torch + output = torch_implementation(X) + + # reduce-overhead compilation may reuse CUDA-graph output buffers. Keep + # each optimizer callback result independent from the next invocation. + if self._compiled: + output = output.clone() + return output + + def __call__(self, G: Tensor, epsilon: float | Tensor = 1e-7, num_experts: int = 1) -> Tensor: + """Orthogonalize XTuner's concatenated regular or MoE matrices. + + Args: + G (Tensor): Concatenated 2D update matrix. + epsilon (float | Tensor): Compatibility argument for Muon's orthogonalization callback. The configured + instance epsilon is used. + num_experts (int): Number of equal row-wise expert matrices concatenated in ``G``. + + Returns: + Tensor: Orthogonalized update with the original shape and dtype. + """ + if G.ndim != 2: + raise ValueError(f"Gram NS expects a 2D tensor, got shape {tuple(G.shape)}") + if num_experts <= 0 or G.size(0) % num_experts: + raise ValueError(f"Invalid num_experts={num_experts} for input shape {tuple(G.shape)}") + + original_shape = G.shape + original_dtype = G.dtype + X = G.view(num_experts, G.size(0) // num_experts, G.size(1)) + output = self._run_selected(X) + return output.to(original_dtype).view(original_shape) diff --git a/xtuner/v1/optim/muon.py b/xtuner/v1/optim/muon.py index 9b8d001910..e31e4885fc 100644 --- a/xtuner/v1/optim/muon.py +++ b/xtuner/v1/optim/muon.py @@ -287,6 +287,13 @@ class Muon(Optimizer): use_triton (bool): Whether to use Triton kernel for Newton-Schulz. Ignored if custom function is provided. newton_schulz_func (Callable | None): Use a custom Newton-Schulz function for orthogonalization. Signature is `func(input: Tensor, epsilon: float, num_experts: int) -> Tensor`. + use_gram_newton_schulz (bool): Use the CuTeDSL Gram Newton-Schulz implementation. Square and + low-aspect-ratio matrices automatically use its standard Newton-Schulz fallback. + gram_ns_epsilon (float): Normalization epsilon for Gram NS and its standard fallback. + gram_restart_iterations (tuple[int, ...]): Zero-based iterations before which Gram NS materializes + the current update and rebuilds the Gram matrix. ``(2,)`` restarts after two iterations. + gram_ns_min_aspect_ratio (float): Gram NS is used only above this matrix aspect ratio. + gram_ns_torch_compile (bool): Compile Gram NS closures using fullgraph/reduce-overhead mode. enable_all2all (bool): Whether to allow the all-to-all communication strategy. Set to False to force the all-gather + reduce-scatter (AGRS) path for all sharded batches. Useful on cluster topologies where all-to-all is unreliable. @@ -313,6 +320,11 @@ def __init__( flatten: bool = False, use_triton: bool = False, newton_schulz_func: Callable | None = None, + use_gram_newton_schulz: bool = False, + gram_ns_epsilon: float = 1e-7, + gram_restart_iterations: tuple[int, ...] = (2,), + gram_ns_min_aspect_ratio: float = 2.0, + gram_ns_torch_compile: bool = True, enable_all2all: bool = True, remainder_strategy: Literal["agrs", "pad_all2all"] = "agrs", muon_split_sizes: dict[Tensor, tuple[int, ...]] | None = None, @@ -330,6 +342,8 @@ def __init__( raise ValueError(f"Invalid remainder_strategy: {remainder_strategy!r}; expected 'agrs' or 'pad_all2all'.") if not enable_all2all and remainder_strategy == "pad_all2all": raise ValueError("remainder_strategy='pad_all2all' requires enable_all2all=True.") + if use_gram_newton_schulz and (use_triton or newton_schulz_func is not None): + raise ValueError("Gram Newton-Schulz cannot be combined with use_triton or newton_schulz_func.") # Default arguments for each param group defaults = dict( @@ -382,6 +396,17 @@ def __init__( if not callable(newton_schulz_func): raise TypeError(f"newton_schulz_func must be a callable function, got {type(newton_schulz_func)}") self._newton_schulz_func = newton_schulz_func + elif use_gram_newton_schulz: + from .gram_newton_schulz import GramNewtonSchulz + + compile_kwargs = {"fullgraph": True, "mode": "reduce-overhead"} if gram_ns_torch_compile else None + self._newton_schulz_func = GramNewtonSchulz( + epsilon=gram_ns_epsilon, + use_kernels=True, + restart_iterations=gram_restart_iterations, + min_aspect_ratio=gram_ns_min_aspect_ratio, + compile_kwargs=compile_kwargs, + ) elif use_triton: from .newton_schulz_triton import newton_schulz_triton