From 7e98a419a764f9c0806f9cfa8ed3d73455af1a8b Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Thu, 3 Sep 2026 16:57:49 +0800 Subject: [PATCH] fix: resolve audited dispatch, kernel, and rollout bugs - re-register the linear family with the operator dispatcher (ASTR_OPS / op_backend / resolve) - fix bf16 gemv misaligned-address faults and element mispairing for offset weights - reject misaligned bf16_swiglu inputs with a clear error and fall back in the backend gate - make the rollout reuse decision, validation, and return atomic under one policy snapshot - add the documented post-scoring rollout version check - derive live+1 under the scheduler lock in optimizer_step via apply_weight_update(None, ...) - reject rollout_max_policy_lag below rollout_interval - 1 at config time - sync gemv stream-test inputs before switching streams; drop dead loader imports --- astrai/config/train_config.py | 11 ++ astrai/extension/backend/linear.py | 134 +++++++++++++++++++++++- astrai/extension/backend/swiglu.py | 5 + astrai/extension/loader.py | 3 - astrai/inference/scheduler.py | 16 ++- astrai/trainer/rollout.py | 75 ++++++++----- astrai/trainer/strategy.py | 5 +- csrc/kernels/bf16_gemv.cu | 19 ++-- csrc/kernels/bf16_swiglu.cu | 20 ++++ tests/extension/test_gemv.py | 38 +++++++ tests/extension/test_linear_dispatch.py | 58 ++++++++++ tests/extension/test_swiglu.py | 13 +++ tests/extension/test_swiglu_dispatch.py | 19 ++++ tests/inference/test_scheduler.py | 26 +++++ tests/trainer/test_online_e2e.py | 35 +++++++ tests/trainer/test_online_strategy.py | 3 + tests/trainer/test_rollout.py | 54 +++++++++- 17 files changed, 485 insertions(+), 49 deletions(-) diff --git a/astrai/config/train_config.py b/astrai/config/train_config.py index 6a62be9..b27bd88 100644 --- a/astrai/config/train_config.py +++ b/astrai/config/train_config.py @@ -234,4 +234,15 @@ class TrainConfig(BaseConfig): f"numbers of forward passes and desynchronize the " f"ddp/fsdp collectives, deadlocking NCCL" ) + if ( + self.rollout_max_policy_lag is not None + and self.rollout_max_policy_lag < self.rollout_interval - 1 + ): + raise ValueError( + f"rollout_max_policy_lag={self.rollout_max_policy_lag} " + f"cannot be below rollout_interval - 1 = " + f"{self.rollout_interval - 1}: the replay cache reuses " + f"rollouts up to that lag, so a tighter bound guarantees " + f"a fatal RolloutVersionError mid-training" + ) return self diff --git a/astrai/extension/backend/linear.py b/astrai/extension/backend/linear.py index 0af5f04..97b580e 100644 --- a/astrai/extension/backend/linear.py +++ b/astrai/extension/backend/linear.py @@ -6,15 +6,31 @@ selection is keyed on the decode batch size alone (M in [2, 4], where it sits at the HBM bandwidth floor and beat the cuBLAS small-M path on every measured family). Every training, prefill-sized, out-of-band, or unsupported call falls back to PyTorch. + +The family stays registered with the shared operator dispatcher, so +``op_backend(linear=...)``, ``ASTR_OPS=linear=...``, and ``resolve`` / +``explain`` keep working like for attention and rotary. The per-layer +hot path only consults the dispatcher when one of those selections is +active, keeping it free of axes dictionaries and record sorting. """ -from typing import Optional +from typing import Any, Dict, List, Optional import torch import torch.nn.functional as F from torch import Tensor -from astrai.extension.dispatch import env_mode +from astrai.extension.dispatch import ( + ImplRecord, + Spec, + axis, + env_mode, + env_selection, + get_override, + register_family, + resolve, + tensor_axes, +) from astrai.extension.loader import is_available from astrai.extension.ops.gemv import bf16_gemv @@ -25,6 +41,10 @@ from astrai.extension.ops.gemv import bf16_gemv _AUTO_GEMV_M = frozenset({2, 3, 4}) +def _torch_linear(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor: + return F.linear(x, weight, bias) + + def _inference_bf16_gemv( x: Tensor, weight: Tensor, bias: Optional[Tensor] = None ) -> Tensor: @@ -64,6 +84,108 @@ def _gemv_capable(x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> bool: ) +def _axes(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Dict[str, Any]: + weight_shape = tuple(weight.shape) + m = 1 if x.ndim == 1 else (x.shape[0] if x.ndim == 2 else None) + supported_m = m is not None and 1 <= m <= 8 + shape_matches = ( + weight.ndim == 2 + and x.ndim in (1, 2) + and bool(x.shape) + and x.shape[-1] == weight_shape[-1] + ) + same_device = x.device == weight.device and ( + bias is None or bias.device == x.device + ) + bias_supported = bias is None or ( + bias.ndim == 1 + and weight.ndim == 2 + and bias.shape[0] == weight_shape[0] + and bias.dtype == torch.bfloat16 + and bias.is_contiguous() + ) + capability = torch.cuda.get_device_capability(x.device) if x.is_cuda else None + return tensor_axes( + x, + mode=env_mode("ASTRAI_GEMV"), + m=m, + supported_m=supported_m, + auto_m=m in _AUTO_GEMV_M, + shape_matches=shape_matches, + same_device=same_device, + weight_dtype=weight.dtype, + x_contiguous=x.is_contiguous(), + weight_contiguous=weight.is_contiguous(), + bias_supported=bias_supported, + capability=capability, + ) + + +_SPEC_CAPABLE = ( + axis("device_cuda").truthy() + & axis("dtype").in_(torch.bfloat16) + & axis("weight_dtype").in_(torch.bfloat16) + & axis("grad_enabled").eq(False) + & axis("supported_m").truthy() + & axis("shape_matches").truthy() + & axis("same_device").truthy() + & axis("x_contiguous").truthy() + & axis("weight_contiguous").truthy() + & axis("bias_supported").truthy() + & Spec.of( + lambda ax: ax.get("capability") is not None and ax.get("capability") >= (8, 0), + "capability>=sm_80", + ) +) + +_SPEC_AUTO = _SPEC_CAPABLE & axis("auto_m").truthy() + + +def _linear_records() -> List[ImplRecord]: + mode = env_mode("ASTRAI_GEMV") + gemv_priority = 0 if mode == "1" else 100 + auto_priority = 0 if mode == "auto" else 90 + torch_priority = 0 if mode == "0" else 50 + return [ + ImplRecord( + family="linear", + name="gemv", + obj=_inference_bf16_gemv, + spec=_SPEC_CAPABLE, + available=lambda: is_available("bf16_gemv"), + priority=gemv_priority, + ), + ImplRecord( + family="linear", + name="auto_gemv", + obj=_inference_bf16_gemv, + spec=_SPEC_AUTO, + available=lambda: is_available("bf16_gemv"), + priority=auto_priority, + ), + ImplRecord( + family="linear", + name="torch", + obj=_torch_linear, + spec=Spec.always(), + priority=torch_priority, + ), + ] + + +def _fallback_record() -> ImplRecord: + return ImplRecord( + family="linear", + name="torch", + obj=_torch_linear, + spec=Spec.always(), + priority=999, + ) + + +register_family("linear", _axes, _linear_records, _fallback_record) + + def linear(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor: """Apply a linear projection with safe inference-only GEMV dispatch. @@ -71,12 +193,18 @@ def linear(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor: primitive can safely handle any M in ``{1, ..., 8}``, and ``auto`` (the default) uses GEMV for decode batches with M in ``{2, 3, 4}``. """ + # Route through the shared dispatcher whenever a selection is active so + # explicit/context/env overrides stay honored; otherwise keep the hot + # path free of axes dictionaries and record sorting. + if get_override("linear") is not None or env_selection("linear") is not None: + return resolve("linear", x, weight, bias).record.obj(x, weight, bias) + mode = env_mode("ASTRAI_GEMV") if mode != "0" and _gemv_capable(x, weight, bias): m = 1 if x.ndim == 1 else x.shape[0] if mode == "1" or m in _AUTO_GEMV_M: return _inference_bf16_gemv(x, weight, bias) - return F.linear(x, weight, bias) + return _torch_linear(x, weight, bias) __all__ = ["linear"] diff --git a/astrai/extension/backend/swiglu.py b/astrai/extension/backend/swiglu.py index 805023c..c838781 100644 --- a/astrai/extension/backend/swiglu.py +++ b/astrai/extension/backend/swiglu.py @@ -40,6 +40,11 @@ def _swiglu_capable(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> bool: or not x.is_contiguous() or not up_weight.is_contiguous() or not gate_weight.is_contiguous() + # The fused kernel reads all streams as uint4; contiguous-but-offset + # views are routed to the unfused chain instead of failing. + or (x.data_ptr() & 15) != 0 + or (up_weight.data_ptr() & 15) != 0 + or (gate_weight.data_ptr() & 15) != 0 or not is_available("bf16_swiglu") ) diff --git a/astrai/extension/loader.py b/astrai/extension/loader.py index b54dbf6..da8f28f 100644 --- a/astrai/extension/loader.py +++ b/astrai/extension/loader.py @@ -20,11 +20,8 @@ import glob import importlib import logging import os -from functools import cache from typing import Dict, List -import torch - logger = logging.getLogger(__name__) _LIB_DIR = os.path.join(os.path.dirname(__file__), "lib") diff --git a/astrai/inference/scheduler.py b/astrai/inference/scheduler.py index 3de9ac0..12e2a1a 100644 --- a/astrai/inference/scheduler.py +++ b/astrai/inference/scheduler.py @@ -176,11 +176,21 @@ class InferenceScheduler: return self._commit_weight_version(policy_version) @_with_weight_lock - def apply_weight_update(self, policy_version: int, update: Callable[[], T]) -> T: - """Mutate shared weights and publish their version without generation.""" + def apply_weight_update( + self, policy_version: Optional[int], update: Callable[[], T] + ) -> T: + """Mutate shared weights and publish their version without generation. + + ``policy_version=None`` derives ``live + 1`` under the same lock, for + callers that only need "advance by one" (e.g. ``optimizer.step()``) + without a read-compute-write race on the current version. + """ if not callable(update): raise TypeError("update must be callable") - self._validate_weight_version(policy_version, require_advance=True) + if policy_version is None: + policy_version = self._policy_version + 1 + else: + self._validate_weight_version(policy_version, require_advance=True) self._ensure_weight_update_ready() result = update() diff --git a/astrai/trainer/rollout.py b/astrai/trainer/rollout.py index 7e81d2c..a012353 100644 --- a/astrai/trainer/rollout.py +++ b/astrai/trainer/rollout.py @@ -147,8 +147,15 @@ class RolloutGenerator: with self._weight_lock: return self.scheduler.update_weights(policy_version) - def apply_weight_update(self, policy_version: int, update: Callable[[], T]) -> T: - """Apply a shared-model mutation at an atomic generation boundary.""" + def apply_weight_update( + self, policy_version: Optional[int], update: Callable[[], T] + ) -> T: + """Apply a shared-model mutation at an atomic generation boundary. + + ``policy_version=None`` lets the scheduler derive ``live + 1`` under + the policy lock, closing the read-compute-write race for callers + that only need to advance by one. + """ with self._weight_lock: return self.scheduler.apply_weight_update(policy_version, update) @@ -424,7 +431,9 @@ class RolloutRunner: """Publish the shared policy's new version to the rollout backend.""" return self.generator.update_weights(policy_version) - def apply_weight_update(self, policy_version: int, update: Callable[[], T]) -> T: + def apply_weight_update( + self, policy_version: Optional[int], update: Callable[[], T] + ) -> T: """Apply a model update and publish its version as one operation.""" return self.generator.apply_weight_update(policy_version, update) @@ -502,35 +511,43 @@ class RolloutRunner: """Return ``(cached or fresh) RolloutResult`` plus an ``is_fresh`` flag. Triggers a new rollout when ``_steps_since_rollout >= rollout_interval`` - or when the cache is empty. + or when the cache is empty. The reuse decision, its version + validation, and the returned object are all captured inside one + policy snapshot, so a concurrent commit, refresh, or cache clear + can never hand out an object the snapshot has already invalidated. """ cache_key = self._batch_key(batch) - if ( - self._cache is None - or cache_key != self._cache_key - or self._steps_since_rollout >= self.rollout_interval - ): - raw = self.generator.generate(batch) - self._validate_policy_version(raw) - scored = self._score(raw) - def commit(live_version: int) -> Tuple[RolloutResult, bool]: - self._validate_policy_version(scored, live_version=live_version) - self._cache = scored - self._cache_key = cache_key - self._steps_since_rollout = 0 - return scored, True - - # A weight update cannot land between the final version check and - # cache publication. Reward scoring itself intentionally remains - # outside the policy lock because it may call an external service. - return self.generator.with_policy_snapshot(commit) - - cached = self._cache - assert cached is not None - - def reuse(live_version: int) -> Tuple[RolloutResult, bool]: + def reuse(live_version: int) -> Optional[Tuple[RolloutResult, bool]]: + cached = self._cache + if ( + cached is None + or self._cache_key != cache_key + or self._steps_since_rollout >= self.rollout_interval + ): + return None self._validate_policy_version(cached, live_version=live_version) return cached, False - return self.generator.with_policy_snapshot(reuse) + outcome = self.generator.with_policy_snapshot(reuse) + if outcome is not None: + return outcome + + raw = self.generator.generate(batch) + self._validate_policy_version(raw) + scored = self._score(raw) + # Post-scoring check: reward scoring may call slow external services; + # surface an over-lag policy move before the commit critical section. + self._validate_policy_version(scored) + + def commit(live_version: int) -> Tuple[RolloutResult, bool]: + self._validate_policy_version(scored, live_version=live_version) + self._cache = scored + self._cache_key = cache_key + self._steps_since_rollout = 0 + return scored, True + + # A weight update cannot land between the final version check and + # cache publication. Reward scoring itself intentionally remains + # outside the policy lock because it may call an external service. + return self.generator.with_policy_snapshot(commit) diff --git a/astrai/trainer/strategy.py b/astrai/trainer/strategy.py index 538fc79..151c7e6 100644 --- a/astrai/trainer/strategy.py +++ b/astrai/trainer/strategy.py @@ -292,8 +292,9 @@ class BaseStrategy(ABC): if self._rollout_runner is None: return optimizer.step() - next_version = self.policy_version + 1 - result = self._rollout_runner.apply_weight_update(next_version, optimizer.step) + # None lets the scheduler derive live+1 under the policy lock, + # avoiding a read-compute-write race on policy_version. + result = self._rollout_runner.apply_weight_update(None, optimizer.step) self._rollout_runner.step() return result diff --git a/csrc/kernels/bf16_gemv.cu b/csrc/kernels/bf16_gemv.cu index 1803d59..e302663 100644 --- a/csrc/kernels/bf16_gemv.cu +++ b/csrc/kernels/bf16_gemv.cu @@ -52,15 +52,18 @@ __global__ void bf16_gemv_kernel( const int wtail_start = whead + wvecs * 8; const uint4* __restrict__ w4 = reinterpret_cast(wrow + whead); - // x chunks pair element-for-element with the aligned weight middle. When - // K % 8 == 0 every x row base shares the weight alignment, so one pure - // uint4 loop covers all rows (the production case: head/tail empty, no - // branching inside the loop). Otherwise per-row uint4 loads are not - // 16-byte addressable, and scalar x pairing keeps the kernel correct for - // any K while the weight stream stays vectorized. + // x chunks pair element-for-element with the aligned weight middle: + // the uint4 view is rooted at ``x + whead`` (16-byte aligned by the + // branch guard), and each row strides by ``k / 8`` vectors because its + // first middle element sits ``whead`` scalars past ``row * k``. When + // K % 8 == 0 and the weight row is already aligned (whead == 0, the + // production case) this reduces to one pure uint4 loop with an empty + // head/tail. Otherwise per-row uint4 loads are not 16-byte addressable, + // and scalar x pairing keeps the kernel correct for any K while the + // weight stream stays vectorized. if (k % 8 == 0 && ((reinterpret_cast(x) + 2u * static_cast(whead)) & 15u) == 0u) { - const auto* x4 = reinterpret_cast(x); + const auto* x4 = reinterpret_cast(x + whead); for (int v = threadIdx.x; v < wvecs; v += blockDim.x) { const uint4 wv_raw = w4[v]; const auto* wv = @@ -68,7 +71,7 @@ __global__ void bf16_gemv_kernel( #pragma unroll for (int row = 0; row < Rows; ++row) { const uint4 xv_raw = - x4[(static_cast(row) * wvecs) + v]; + x4[(static_cast(row) * (k / 8)) + v]; const auto* xv = reinterpret_cast(&xv_raw); #pragma unroll diff --git a/csrc/kernels/bf16_swiglu.cu b/csrc/kernels/bf16_swiglu.cu index 749d768..1ab4ff2 100644 --- a/csrc/kernels/bf16_swiglu.cu +++ b/csrc/kernels/bf16_swiglu.cu @@ -170,6 +170,26 @@ torch::Tensor bf16_swiglu( gate_weight.is_contiguous(), "x and weights must be contiguous" ); + // The kernel loads all three streams as uint4; contiguous-but-offset + // views would fault with an opaque "misaligned address" CUDA error, so + // reject them here with an actionable message. + TORCH_CHECK( + (reinterpret_cast(x.data_ptr()) & 15u) == 0u, + "bf16_swiglu requires 16-byte aligned x (storage_offset must keep " + "data_ptr divisible by 16); clone the tensor or use the torch path" + ); + TORCH_CHECK( + (reinterpret_cast(up_weight.data_ptr()) & 15u) == 0u, + "bf16_swiglu requires 16-byte aligned up_weight (storage_offset " + "must keep data_ptr divisible by 16); clone the tensor or use the " + "torch path" + ); + TORCH_CHECK( + (reinterpret_cast(gate_weight.data_ptr()) & 15u) == 0u, + "bf16_swiglu requires 16-byte aligned gate_weight (storage_offset " + "must keep data_ptr divisible by 16); clone the tensor or use the " + "torch path" + ); TORCH_CHECK( !x.requires_grad() && !up_weight.requires_grad() && !gate_weight.requires_grad(), diff --git a/tests/extension/test_gemv.py b/tests/extension/test_gemv.py index 8aa3fd2..4432f27 100644 --- a/tests/extension/test_gemv.py +++ b/tests/extension/test_gemv.py @@ -130,6 +130,7 @@ def test_bf16_gemv_small_batch_fuses_bias(): def test_bf16_gemv_uses_current_stream(): x = torch.randn(1536, device="cuda", dtype=torch.bfloat16) weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16) + torch.cuda.synchronize() stream = torch.cuda.Stream() with torch.cuda.stream(stream): actual = bf16_gemv(x, weight) @@ -171,6 +172,43 @@ def test_bf16_gemv_handles_unaligned_k(n, k): torch.testing.assert_close(actual3, F.linear(x3, weight), rtol=0.02, atol=0.5) +@skip_no_gemv +@pytest.mark.parametrize("m", [1, 2, 3, 4]) +def test_bf16_gemv_accepts_complementary_misalignment(m): + """Misaligned weight rows plus an x base chosen so the vectorized branch + is entered with a non-16B-aligned ``x`` pointer (regression: the branch + guard checked ``x + whead`` alignment but the uint4 view was rooted at + ``x`` itself, faulting with a misaligned-address CUDA error).""" + torch.manual_seed(37) + n, k = 256, 1536 + # offset 5 elements = +10 bytes: weight rows land at 10 % 16 (whead=3) + # and x at 10 % 16, so (x + 2*whead) % 16 == 0 selects the fast path. + big_w = torch.randn(n * k + 8, device="cuda", dtype=torch.bfloat16) + weight = big_w[5 : 5 + n * k].view(n, k) + big_x = torch.randn(m * k + 8, device="cuda", dtype=torch.bfloat16) + x = big_x[5 : 5 + m * k].view(m, k) if m > 1 else big_x[5 : 5 + k] + assert (x.data_ptr() & 15) == 10 and (weight.data_ptr() & 15) == 10 + + actual = bf16_gemv(x, weight) + expected = F.linear(x, weight) + torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.5 if m > 1 else 0.25) + + +@skip_no_gemv +def test_bf16_gemv_scalar_path_handles_misaligned_weight_only(): + """Weight rows misaligned while x stays 16B-aligned take the scalar-x + middle and must stay exact.""" + torch.manual_seed(41) + n, k = 256, 1536 + big_w = torch.randn(n * k + 8, device="cuda", dtype=torch.bfloat16) + weight = big_w[5 : 5 + n * k].view(n, k) + x = torch.randn(2, k, device="cuda", dtype=torch.bfloat16) + assert (weight.data_ptr() & 15) == 10 and (x.data_ptr() & 15) == 0 + + actual = bf16_gemv(x, weight) + torch.testing.assert_close(actual, F.linear(x, weight), rtol=0.02, atol=0.5) + + @skip_no_gemv def test_bf16_gemv_small_batch_cuda_graph_replay(): torch.manual_seed(31) diff --git a/tests/extension/test_linear_dispatch.py b/tests/extension/test_linear_dispatch.py index ed00093..56c0adf 100644 --- a/tests/extension/test_linear_dispatch.py +++ b/tests/extension/test_linear_dispatch.py @@ -7,6 +7,7 @@ import torch.nn.functional as F from astrai.extension import is_available, linear from astrai.extension.backend import linear as public_linear +from astrai.extension.dispatch import explain, op_backend, resolve # The package attribute ``linear`` is the dispatched function; reach the # module object explicitly for monkeypatching its private helpers. @@ -166,3 +167,60 @@ def test_dispatched_linear_cuda_graph_replay(monkeypatch): graph.replay() expected = F.linear(x, weight) torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25) + + +def test_linear_family_is_registered_with_shared_dispatcher(): + from astrai.extension.dispatch import _FAMILIES + + assert "linear" in _FAMILIES + x = torch.randn(2, 8) + weight = torch.randn(4, 8) + resolution = resolve("linear", x, weight) + assert resolution.record.family == "linear" + assert resolution.origin in ("chain", "fallback") + assert "linear" in explain("linear", x, weight) + + +@skip_no_gemv +def test_ops_env_override_forces_torch_for_capable_call(monkeypatch): + """ASTR_OPS=linear=torch must keep working after the M-band rewrite + (regression: the family was silently dropped from the dispatcher, so + the override warned, fell through, and the gemv kernel still ran).""" + monkeypatch.setenv("ASTR_OPS", "linear=torch") + x = torch.randn(2, 1536, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + assert not _routes_to_gemv(monkeypatch, x, weight) + torch.testing.assert_close( + linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25 + ) + + +@skip_no_gemv +def test_ops_env_override_forces_gemv(monkeypatch): + monkeypatch.setenv("ASTR_OPS", "linear=gemv") + x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + # M=1 is outside the auto band but inside the forced gemv record. + assert _routes_to_gemv(monkeypatch, x, weight) + + +@skip_no_gemv +def test_op_backend_context_selects_torch(monkeypatch): + monkeypatch.setenv("ASTRAI_GEMV", "1") + x = torch.randn(2, 1536, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(), op_backend(linear="torch"): + assert not _routes_to_gemv(monkeypatch, x, weight) + torch.testing.assert_close( + linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25 + ) + # The override is scoped: the forced mode applies again afterwards. + with torch.no_grad(): + assert _routes_to_gemv(monkeypatch, x, weight) + + +def test_op_backend_rejects_unknown_linear_handle(): + with pytest.raises(ValueError, match="Unknown linear implementation"): + op_backend(linear="nonexistent").__enter__() diff --git a/tests/extension/test_swiglu.py b/tests/extension/test_swiglu.py index 04165b7..f5b1c2d 100644 --- a/tests/extension/test_swiglu.py +++ b/tests/extension/test_swiglu.py @@ -97,3 +97,16 @@ def test_bf16_swiglu_uses_current_stream_and_cuda_graph(): def test_bf16_swiglu_rejects_unsupported_inputs(make_args, error): with pytest.raises(RuntimeError, match=error): bf16_swiglu(*make_args()) + + +@skip_no_swiglu +def test_bf16_swiglu_rejects_misaligned_storage_with_clear_error(): + """Contiguous-but-offset views must fail the wrapper's TORCH_CHECK with + an actionable message instead of a sticky CUDA misaligned-address error + (regression: the kernel casts x directly to uint4 without checking).""" + k = 1536 + base = torch.randn(k + 1, device="cuda", dtype=torch.bfloat16) + x = base[1:] # +2 bytes: contiguous but not 16B-aligned + weights = torch.randn(8, k, device="cuda", dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match="16-byte aligned"): + bf16_swiglu(x, weights, weights) diff --git a/tests/extension/test_swiglu_dispatch.py b/tests/extension/test_swiglu_dispatch.py index 130e5a7..4611242 100644 --- a/tests/extension/test_swiglu_dispatch.py +++ b/tests/extension/test_swiglu_dispatch.py @@ -95,3 +95,22 @@ def test_auto_uses_unfused_chain_until_shape_is_qualified(monkeypatch): actual = swiglu(x, up_weight, gate_weight) expected = reference_swiglu(x, up_weight, gate_weight) torch.testing.assert_close(actual, expected, rtol=0.03, atol=0.1) + + +@skip_no_swiglu +def test_mode_one_falls_back_for_misaligned_storage(monkeypatch): + """Contiguous-but-offset views must route to the unfused torch chain + even with ASTRAI_SWIGLU=1 instead of reaching the uint4-only kernel + (regression: the fused primitive faulted with a misaligned-address + CUDA error for such inputs).""" + monkeypatch.setenv("ASTRAI_SWIGLU", "1") + k = 1536 + x_base = torch.randn(2 * k + 8, device="cuda", dtype=torch.bfloat16) * 0.1 + x = x_base[1 : 1 + 2 * k].view(2, k) + assert x.is_contiguous() and (x.data_ptr() & 15) != 0 + up_weight = torch.randn(64, k, device="cuda", dtype=torch.bfloat16) * 0.02 + gate_weight = torch.randn_like(up_weight) + with torch.no_grad(): + actual = swiglu(x, up_weight, gate_weight) + expected = reference_swiglu(x, up_weight, gate_weight) + torch.testing.assert_close(actual, expected, rtol=0.03, atol=0.01) diff --git a/tests/inference/test_scheduler.py b/tests/inference/test_scheduler.py index 292b3cd..82d7496 100644 --- a/tests/inference/test_scheduler.py +++ b/tests/inference/test_scheduler.py @@ -583,6 +583,32 @@ def test_scheduler_applies_weight_mutation_and_version_atomically(device): with pytest.raises(RuntimeError, match="optimizer failed"): scheduler.apply_weight_update(2, failed_mutation) assert scheduler.policy_version == 1 + + # None derives live+1 under the lock: no read-compute-write race + # on the current version for advance-by-one callers. + assert scheduler.apply_weight_update(None, mutate) == "updated" + assert scheduler.policy_version == 2 + finally: + scheduler.stop() + + +def test_scheduler_atomic_advance_survives_interleaved_publish(device): + """A concurrent publish between reading the live version and applying + the update must not fail ``require_advance`` (regression: callers + computed live+1 outside the lock, a TOCTOU that raised spuriously).""" + scheduler, _tok, _model = _make_real_scheduler(device) + + try: + # Simulate the race directly: a version read that goes stale before + # apply_weight_update acquires the lock. With None the scheduler + # re-derives live+1 inside the critical section. + stale_read = scheduler.policy_version + 1 + scheduler.update_weights(1) + assert stale_read == 1 # now equals live -> explicit form would raise + with pytest.raises(ValueError, match="must advance"): + scheduler.apply_weight_update(stale_read, lambda: "ok") + assert scheduler.apply_weight_update(None, lambda: "ok") == "ok" + assert scheduler.policy_version == 2 finally: scheduler.stop() diff --git a/tests/trainer/test_online_e2e.py b/tests/trainer/test_online_e2e.py index 9e0aaca..e88e650 100644 --- a/tests/trainer/test_online_e2e.py +++ b/tests/trainer/test_online_e2e.py @@ -144,3 +144,38 @@ def test_online_rollout_end_to_end( checkpoint = Checkpoint.load(checkpoint_dir) assert checkpoint.meta["policy_version"] == 2 assert len(created_reference_models) == 1 + + +def _minimal_online_config(**overrides): + """A TrainConfig for online GRPO that only needs field overrides.""" + defaults = dict( + strategy="online_grpo", + model_fn=lambda: torch.nn.Linear(2, 2), + dataset=InstructionDataset(), + optimizer_fn=lambda m: torch.optim.SGD(m.parameters(), lr=0.0), + scheduler_fn=lambda o: SchedulerFactory.create( + "cosine", o, warmup_steps=1, lr_decay_steps=4, min_rate=0.05 + ), + reward_model_fn=LengthRewardModel, + ) + defaults.update(overrides) + return TrainConfig(**defaults) + + +def test_online_config_rejects_contradictory_policy_lag(): + """rollout_max_policy_lag below rollout_interval - 1 guarantees a fatal + RolloutVersionError mid-training; it must fail at config time instead.""" + with pytest.raises(ValueError, match="rollout_max_policy_lag=0"): + _minimal_online_config(rollout_interval=3, rollout_max_policy_lag=0) + + # lag == interval - 1 (including the derived default) stays valid. + config = _minimal_online_config(rollout_interval=3, rollout_max_policy_lag=2) + assert config.rollout_max_policy_lag == 2 + config = _minimal_online_config(rollout_interval=3) + assert config.rollout_max_policy_lag is None + + # Offline strategies never consult the rollout window. + config = _minimal_online_config( + strategy="sft", rollout_interval=3, rollout_max_policy_lag=0 + ) + assert config.rollout_max_policy_lag == 0 diff --git a/tests/trainer/test_online_strategy.py b/tests/trainer/test_online_strategy.py index 396231f..dbb2d0f 100644 --- a/tests/trainer/test_online_strategy.py +++ b/tests/trainer/test_online_strategy.py @@ -62,6 +62,9 @@ class _RecordingRunner: def apply_weight_update(self, policy_version, update): result = update() + if policy_version is None: + # Mirror the scheduler: None derives live+1 under the lock. + policy_version = self.policy_version + 1 self.update_weights(policy_version) return result diff --git a/tests/trainer/test_rollout.py b/tests/trainer/test_rollout.py index e22c04a..385e368 100644 --- a/tests/trainer/test_rollout.py +++ b/tests/trainer/test_rollout.py @@ -459,7 +459,7 @@ def test_rollout_runner_publishes_cache_before_concurrent_policy_update(device): nonlocal validation_calls validation_calls += 1 original_validate(result, live_version=live_version) - if validation_calls == 2: + if validation_calls == 3: final_validation_started.set() assert allow_final_validation_to_finish.wait(timeout=5) @@ -503,6 +503,58 @@ def test_rollout_runner_derives_default_policy_lag_from_interval(device): assert runner.max_policy_lag == 3 +def _interleave_before_snapshot(runner, callback): + """Wrap ``with_policy_snapshot`` so ``callback`` runs just before a + named snapshot callback enters the generator/scheduler locks.""" + original_snapshot = runner.generator.with_policy_snapshot + + def wrapper(inspect): + if inspect.__name__ == "reuse": + callback() + return original_snapshot(inspect) + + runner.generator.with_policy_snapshot = wrapper + + +def test_rollout_runner_reuse_reads_cache_inside_the_snapshot(device): + """The reuse decision must observe the cache under the policy snapshot + (regression: the cache was read outside the lock, so a concurrent + commit between the read and the lock silently handed the trainer a + stale rollout — a lost update).""" + import dataclasses + + runner, _ = _make_runner(device, rollout_interval=100) + batch = _make_instruction_batch(n=1) + first, _ = runner(batch) + assert first.policy_version == 0 + + def concurrent_refresh(): + runner.update_weights(1) + runner._cache = dataclasses.replace(first, policy_version=1) + runner._steps_since_rollout = 0 + + _interleave_before_snapshot(runner, concurrent_refresh) + result, fresh = runner(batch) + assert fresh is False + assert result is not first + assert result.policy_version == 1 + + +def test_rollout_runner_recovers_when_cache_cleared_before_reuse_snapshot(device): + """A cache clear between the reuse decision and the snapshot must + trigger a fresh rollout instead of an assertion failure (regression: + ``assert cached is not None`` fired because the object was captured + outside the lock).""" + runner, _ = _make_runner(device, rollout_interval=100) + batch = _make_instruction_batch(n=1) + first, _ = runner(batch) + + _interleave_before_snapshot(runner, runner.clear_cache) + result, fresh = runner(batch) + assert fresh is True + assert result is not first + + @pytest.mark.parametrize( ("kwargs", "message"), [