chore: 版本号升至 1.3.5

- get_rotary_emb() 返回复数张量替代 Tuple[cos, sin]
- RotaryEmbedding 存储单一 freqs_cis buffer 替代分离的 cos_cached/sin_cached
- forward 中 view_as_complex 重建复数

README.md

@@ -65,6 +65,16 @@ For development dependencies:
 pip install -e ".[dev]"
 ```
 
+#### Download Pre-trained Model
+
+Download pre-trained model weights (1B bilingual checkpoint) to `params/`:
+
+```bash
+python scripts/demo/download.py
+```
+
+Or download manually from [HuggingFace](https://huggingface.co/ViperEk/KHAOSZ) into `params/`.
+
 #### Train a Model
 
 ```bash

assets/docs/README-zh-CN.md

@@ -71,6 +71,16 @@ pip install -e .
 pip install -e ".[dev]"
 ```
 
+#### 下载预训练模型
+
+下载预训练模型权重（1B 双语检查点）到 `params/` 目录：
+
+```bash
+python scripts/demo/download.py
+```
+
+或从 [HuggingFace](https://huggingface.co/ViperEk/KHAOSZ) 手动下载放入 `params/`。
+
 #### 训练模型
 
 ```bash

assets/docs/dataflow.md

@@ -88,7 +88,7 @@ flowchart LR
 - **`DecoderBlock`**: GQA attention + residual + MLP + RMSNorm
 - **`GQA`**: Grouped Query Attention (also `MLA` for multi-latent attention)
 - **`MLP`**: `SiLU(gate(x)) * up(x)` → down projection
-- **`RotaryEmbedding`**: RoPE cos/sin cache
+- **`RotaryEmbedding`**: RoPE complex cache (freqs_cis)
 - **`RMSNorm`**: Layer normalization
 
 ### 4. Training Module
@@ -104,22 +104,23 @@ The training loop is nested: **epoch** → **batch** (with step phase interspers
 ```
 on_train_begin
   on_epoch_begin
-    for each batch:
+    for each accumulation window of batches:        ← step phase
-      if iteration % accumulation_steps == 0:        ← step phase
+      on_step_begin
-        on_step_begin → optimizer.step() → zero_grad → on_step_end
+        for each batch in window:                    ← batch phase
-                                                      ← batch phase
+          on_batch_begin → strategy(batch) → loss → backward → on_batch_end
-      on_batch_begin → strategy(batch) → loss → backward → on_batch_end
+          iteration += 1
-      iteration += 1
+      on_step_end
+      optimizer.step() → zero_grad
 
     on_epoch_end
 on_train_end
 ```
 
 Key points:
-- `on_step_*` wraps optimizer step (fires every `accumulation_steps` batches)
+- `on_step_*` fires every `accumulation_steps` batches, wrapping optimizer step AFTER the hook
-- `on_batch_*` wraps loss computation (fires every batch)
+- `on_batch_*` fires every batch, wrapping loss computation
-- `SchedulerCallback` fires on `on_batch_end` — LR scheduler steps every batch
+- `GradientClippingCallback` fires on `on_step_end`
-- `GradientClippingCallback` fires on `on_step_begin`
+- LR scheduler steps inline (no `SchedulerCallback` class)
 
 #### 4.3 Strategy (`strategy.py`)
 - **`SEQStrategy`**: Next-token prediction, cross-entropy with label smoothing
@@ -136,8 +137,7 @@ Key points:
 - **`CheckpointCallback`**: Saves safetensors at `ckpt_interval` iterations
 - **`ProgressBarCallback`**: tqdm progress display
 - **`MetricLoggerCallback`**: Writes JSONL metrics to `{ckpt_dir}/logs/`
-- **`GradientClippingCallback`**: `clip_grad_norm_` on `on_step_begin`
+- **`GradientClippingCallback`**: `clip_grad_norm_` on `on_step_end`
-- **`SchedulerCallback`**: `scheduler.step()` on `on_batch_end`
 
 ### 5. Inference Module
 

assets/docs/design.md

@@ -91,8 +91,8 @@ classDiagram
         }
 
         class BaseStorage {
-            +Dict segments
+            +MultiSegmentFetcher _fetcher
-            +List keys
+            +keys (property)
             +load(load_path, tokenizer)
             +fetch(begin, end, keys)
             +__len__()
@@ -145,7 +145,7 @@ classDiagram
             +ModelConfig config
             +Registry _registry
             +register(model_type) decorator
-            +get_model_class(model_type) Type
+            +get_component_class(model_type) Type
             +from_pretrained(path, disable_random_init) nn.Module
             +save_pretrained(save_directory)
             +to(*args, **kwargs) Self
@@ -214,7 +214,7 @@ classDiagram
             +int dim
             +int max_len
             +float base
-            +forward(x, position_ids=None) Tuple[Tensor, Tensor]
+            +forward(x, position_ids=None) Tensor
         }
 
         class Embedding {
@@ -225,13 +225,10 @@ classDiagram
 
     namespace tokenize {
         class AutoTokenizer {
-            +List[int] stop_ids
-            +int bos_id
-            +int eos_id
-            +int pad_id
             +vocab_size int
             +encode(tokens, out_ids, add_special_tokens) List[int]
             +decode(tokens, skip_special_tokens) str
+            +__getattr__(name) Any (bos_id, eos_id, pad_id, stop_ids)
             +apply_chat_template(messages, tokenize) Union[str, List[int]]
             +set_chat_template(template)
             +load(path)
@@ -325,6 +322,8 @@ classDiagram
             +float clip_eps
             +float kl_coef
             +int group_size
+            +str reduction
+            +int sync_interval
             +compute_loss(batch) Tensor
         }
 
@@ -369,11 +368,6 @@ classDiagram
             +on_step_begin(context)
         }
 
-        class SchedulerCallback {
-            +on_train_begin(context)
-            +on_batch_end(context)
-        }
-
         class CheckpointCallback {
             +str save_dir
             +int interval
@@ -409,8 +403,6 @@ classDiagram
             +nn.Module model
             +AutoTokenizer tokenizer
             +InferenceScheduler scheduler
-            +int max_batch_size
-            +Optional int max_seq_len
             +generate(prompt, stream, max_tokens, temperature, top_p, top_k) Union[Generator, str, List[str]]
             +generate_with_request(request) Union[Generator, str, List[str]]
             +generate_async(prompt, max_tokens, temperature, top_p, top_k) AsyncGenerator
@@ -421,13 +413,12 @@ classDiagram
         class InferenceScheduler {
             +nn.Module model
             +AutoTokenizer tokenizer
-            +KVCache page_cache
+            +KVCache _page_cache
             +int max_batch_size
             +int max_seq_len
             +int max_prompt_len
             +int page_size
-            +List waiting_queue
+            +TaskManager _task_mgr
-            +List active_tasks
             +add_task(prompt, max_tokens, temperature, top_p, top_k, stream_callback) str
             +remove_task(task_id)
             +start()
@@ -568,7 +559,7 @@ classDiagram
         }
 
         class GenerateResult {
-            +List[str] tokens
+            +List[Tuple[int, str]] tokens
             +List[str] results
             +List[bool] _done
             +append(token, idx)
@@ -643,7 +634,6 @@ classDiagram
     BaseScheduler <|-- SGDRScheduler
     CallbackFactory ..> TrainCallback : creates
     TrainCallback <|-- GradientClippingCallback
-    TrainCallback <|-- SchedulerCallback
     TrainCallback <|-- CheckpointCallback
     TrainCallback <|-- ProgressBarCallback
     TrainCallback <|-- MetricLoggerCallback

astrai/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.3.4"
+__version__ = "1.3.5"
 __author__ = "ViperEkura"
 
 from astrai.config import (

astrai/inference/core/cache.py

@@ -235,7 +235,16 @@ class Storage:
             write_end = min(page_start + page_size, start_pos + seq_len)
             offset = write_start - page_start
             chunk = write_end - write_start
-            if (phys_pages < 0).any():
+            valid = phys_pages >= 0
+            if not valid.all():
+                if valid.any():
+                    valid_pages = phys_pages[valid]
+                    self.k_cache[layer_id, valid_pages, offset : offset + chunk] = k[
+                        valid, written : written + chunk
+                    ]
+                    self.v_cache[layer_id, valid_pages, offset : offset + chunk] = v[
+                        valid, written : written + chunk
+                    ]
                 written += chunk
                 continue
             self.k_cache[layer_id, phys_pages, offset : offset + chunk] = k[
@@ -254,6 +263,16 @@ class Storage:
         v = self.v_cache[layer_id, safe]
         k = k.flatten(1, 2)
         v = v.flatten(1, 2)
+        if (page_table < 0).any():
+            invalid = (
+                (page_table < 0)
+                .unsqueeze(-1)
+                .expand(-1, -1, self.page_size)
+                .flatten(1, 2)
+            )
+            invalid = invalid[:, :, None, None].expand_as(k)
+            k = k.masked_fill(invalid, 0.0)
+            v = v.masked_fill(invalid, 0.0)
         k = k[:, :total_len]
         v = v[:, :total_len]
         return k, v

astrai/inference/core/executor.py

@@ -38,13 +38,11 @@ class Executor:
         tasks = sorted(tasks, key=lambda t: t.task_id)
         batch_sz = len(tasks)
 
-        seq_len = prompt_len - start_pos
+        input_ids = torch.tensor(
-        input_ids = torch.empty(batch_sz, seq_len, dtype=torch.long, device=self.device)
+            [t.prompt_ids[start_pos:prompt_len] for t in tasks],
-
+            dtype=torch.long,
-        for i, t in enumerate(tasks):
+            device=self.device,
-            input_ids[i] = torch.tensor(
+        )
-                t.prompt_ids[start_pos:prompt_len], device=self.device
-            )
 
         task_ids = [t.task_id for t in tasks]
         page_tables = self.page_cache.make_table_tensor(task_ids, self.device)

astrai/inference/core/scheduler.py

@@ -126,9 +126,7 @@ class InferenceScheduler:
 
                 pos_groups: Dict[int, List[Task]] = {}
                 for t in self._task_mgr.get_active_tasks():
-                    chunk = t.next_pos // self._page_cache.page_size
+                    pos_groups.setdefault(t.next_pos, []).append(t)
-                    key = chunk if chunk <= 1 else 1 << (chunk.bit_length() - 1)
-                    pos_groups.setdefault(key, []).append(t)
 
                 if pos_groups:
                     best_key = max(pos_groups, key=lambda k: len(pos_groups[k]))

astrai/inference/core/task.py

@@ -2,8 +2,9 @@ import logging
 import threading
 import time
 import uuid
+from collections import deque
 from enum import Enum
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, Callable, Deque, Dict, List, Optional
 
 from astrai.tokenize.tokenizer import AutoTokenizer
 
@@ -76,7 +77,7 @@ class TaskManager:
         self.max_seq_len = max_seq_len
         self.max_prompt_len = max_prompt_len
 
-        self.waiting_queue: List[Task] = []
+        self.waiting_queue: Deque[Task] = deque()
         self.active_tasks: List[Task] = []
 
         self._task_event = threading.Event()
@@ -129,7 +130,9 @@ class TaskManager:
     def remove_task(self, task_id: str) -> List[Task]:
         with self._lock:
             removed_active = [t for t in self.active_tasks if t.task_id == task_id]
-            self.waiting_queue = [t for t in self.waiting_queue if t.task_id != task_id]
+            self.waiting_queue = deque(
+                t for t in self.waiting_queue if t.task_id != task_id
+            )
             self.active_tasks = [t for t in self.active_tasks if t.task_id != task_id]
         return removed_active
 
@@ -166,16 +169,18 @@ class TaskManager:
         with self._lock:
             take = min(n, len(self.waiting_queue))
             for _ in range(take):
-                to_add.append(self.waiting_queue.pop(0))
+                to_add.append(self.waiting_queue.popleft())
         return to_add
 
     def activate(self, task: Task) -> None:
         task.status = TaskStatus.RUNNING
-        self.active_tasks.append(task)
+        with self._lock:
+            self.active_tasks.append(task)
 
     def return_to_waiting(self, tasks: List[Task]) -> None:
         with self._lock:
-            self.waiting_queue[:0] = tasks
+            for task in reversed(tasks):
+                self.waiting_queue.appendleft(task)
 
     def has_work(self) -> bool:
         return bool(self.active_tasks or self.waiting_queue)

astrai/inference/engine.py

@@ -59,9 +59,15 @@ class GenerateResult:
     def wait(self, timeout: Optional[float] = None) -> bool:
         return self._event.wait(timeout=timeout)
 
-    def wait_completion(self) -> None:
+    def wait_completion(self, timeout: float = 300.0) -> None:
         with self._cond:
-            self._cond.wait_for(lambda: self._completed >= self._total)
+            if not self._cond.wait_for(
+                lambda: self._completed >= self._total, timeout=timeout
+            ):
+                raise TimeoutError(
+                    f"Generation timeout after {timeout}s "
+                    f"({self._completed}/{self._total} completed)"
+                )
 
     def get_results(self) -> List[str]:
         with self._cond:
@@ -267,7 +273,12 @@ class InferenceEngine:
             prompts, max_tokens, temperature, top_p, top_k
         )
 
-        result.wait_completion()
+        try:
+            result.wait_completion()
+        except TimeoutError:
+            for tid in task_ids:
+                self.scheduler.remove_task(tid)
+            raise
 
         for tid in task_ids:
             self.scheduler.remove_task(tid)

astrai/inference/sample.py

@@ -64,16 +64,26 @@ class TopKStrategy(BaseSamplingStrategy):
     def apply(self, logits, filter_value=-float("inf")):
         tk = self.top_k
         if isinstance(tk, Tensor):
+            tk = tk.to(logits.device, non_blocking=True).long().clamp(min=0)
             max_k = int(tk.max().item())
             if max_k <= 0:
                 return logits
-            k = min(max_k, logits.size(-1))
+            max_k = min(max_k, logits.size(-1))
-        elif tk > 0:
+            values, _ = torch.topk(logits, max_k, dim=-1)
-            k = min(tk, logits.size(-1))
+            per_row_k = tk.clamp(max=max_k)
-        else:
+            thresholds = torch.full_like(logits[..., -1:], -float("inf"))
+            positive = per_row_k > 0
+            if positive.any():
+                row_idx = torch.arange(logits.size(0), device=logits.device)[positive]
+                thresholds[positive] = values[
+                    row_idx, per_row_k[positive] - 1
+                ].unsqueeze(-1)
+            logits[logits < thresholds] = filter_value
             return logits
-        thresholds = torch.topk(logits, k, dim=-1)[0][..., -1:]
+        if tk > 0:
-        logits[logits < thresholds] = filter_value
+            k = min(tk, logits.size(-1))
+            thresholds = torch.topk(logits, k, dim=-1)[0][..., -1:]
+            logits[logits < thresholds] = filter_value
         return logits
 
 

astrai/model/module.py

@@ -1,4 +1,4 @@
-from typing import Optional, Tuple
+from typing import Optional
 
 import torch
 import torch.nn as nn
@@ -25,11 +25,13 @@ def get_rotary_emb(
     max_len: int,
     base: float = 10000,
     device: Optional[torch.device] = None,
-) -> Tuple[Tensor, Tensor]:
+) -> Tensor:
     theta = base ** (-torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim)
     t = torch.arange(0, max_len, dtype=torch.float64, device=device)
-    freqs = torch.outer(t, theta)
+    freqs = torch.outer(t, theta).float()
-    return torch.cos(freqs).float(), torch.sin(freqs).float()
+    cos = torch.cos(freqs)
+    sin = torch.sin(freqs)
+    return torch.complex(cos, sin)
 
 
 def apply_rotary_emb(x: torch.Tensor, freqs_cis: Tensor) -> Tensor:
@@ -50,10 +52,10 @@ class RotaryEmbedding(nn.Module):
         self.base = base
         self._set_rotary_buffer(self.max_len)
 
-    def _set_rotary_buffer(self, max_len: int, device: Optional[torch.device] = None):
+    def _set_rotary_buffer(self, max_len: int):
-        cos_cached, sin_cached = get_rotary_emb(self.dim, max_len, self.base, device)
+        rotary_emb = get_rotary_emb(self.dim, max_len, self.base)
-        self.register_buffer("cos_cached", cos_cached, persistent=False)
+        freqs_cis = torch.view_as_real(rotary_emb)
-        self.register_buffer("sin_cached", sin_cached, persistent=False)
+        self.register_buffer("freqs_cis", freqs_cis, persistent=False)
 
     def forward(self, x: Tensor, position_ids: Optional[Tensor] = None) -> Tensor:
         if position_ids is None:
@@ -62,9 +64,8 @@ class RotaryEmbedding(nn.Module):
                 .unsqueeze(0)
                 .expand(x.size(0), -1)
             )
-        cos = self.cos_cached[position_ids].float()
+        position_freq_cis = self.freqs_cis[position_ids].float()
-        sin = self.sin_cached[position_ids].float()
+        return torch.view_as_complex(position_freq_cis)
-        return torch.complex(cos, sin)
 
 
 class Linear(nn.Module):

astrai/trainer/metric_util.py

@@ -1,75 +1,42 @@
-from typing import Dict
+from typing import Any, Callable, Dict
 
+import torch
 import torch.nn as nn
 
 
-def grad_norm(model: nn.Module, norm_type: int = 2) -> Dict[str, float]:
+def _grad_stat(
-    """Compute gradient norm for each parameter in the model."""
+    model: nn.Module, fn: Callable[[torch.Tensor], Any], default: Any
-    norms = {}
+) -> dict:
+    results = {}
     for name, param in model.named_parameters():
-        norms[name] = 0.0
+        results[name] = default
-        if param.grad:
+        if param.grad is not None:
-            norm = param.grad.data.norm(norm_type).item()
+            results[name] = fn(param.grad.data)
-            norms[name] = norm
+    return results
-    return norms
+
+
+def grad_norm(model: nn.Module, norm_type: int = 2) -> Dict[str, float]:
+    return _grad_stat(model, lambda g: g.norm(norm_type).item(), 0.0)
 
 
 def grad_std(model: nn.Module) -> Dict[str, float]:
-    """Compute standard deviation of gradients for each parameter."""
+    return _grad_stat(model, lambda g: g.std().item(), 0.0)
-    stds = {}
-    for name, param in model.named_parameters():
-        stds[name] = 0.0
-        if param.grad:
-            std = param.grad.data.std().item()
-            stds[name] = std
-    return stds
 
 
 def grad_max(model: nn.Module) -> Dict[str, float]:
-    """Find the maximum absolute gradient value for each parameter."""
+    return _grad_stat(model, lambda g: g.max().item(), -float("inf"))
-    max_vals = {}
-    for name, param in model.named_parameters():
-        max_vals[name] = -float("inf")
-        if param.grad:
-            max_val = param.grad.data.max().item()
-            max_vals[name] = max_val
-
-    return max_vals
 
 
 def grad_min(model: nn.Module) -> Dict[str, float]:
-    """Find the minimum absolute gradient value for each parameter."""
+    return _grad_stat(model, lambda g: g.min().item(), float("inf"))
-    min_vals = {}
-    for name, param in model.named_parameters():
-        min_vals[name] = float("inf")
-        if param.grad:
-            min_val = param.grad.data.min().item()
-            min_vals[name] = min_val
-
-    return min_vals
 
 
 def grad_mean(model: nn.Module) -> Dict[str, float]:
-    """Compute mean of gradients for each parameter."""
+    return _grad_stat(model, lambda g: g.mean().item(), 0.0)
-    means = {}
-    for name, param in model.named_parameters():
-        means[name] = 0.0
-        if param.grad:
-            mean = param.grad.data.mean().item()
-            means[name] = mean
-
-    return means
 
 
 def grad_nan_num(model: nn.Module) -> Dict[str, int]:
-    """Count the number of NaNs in gradients for each parameter."""
+    return _grad_stat(model, lambda g: g.isnan().sum().item(), 0)
-    nan_nums = {}
-    for name, param in model.named_parameters():
-        nan_nums[name] = 0
-        if param.grad:
-            nan_num = param.grad.isnan().sum().item()
-            nan_nums[name] = nan_num
-    return nan_nums
 
 
 def ctx_get_loss(ctx):

astrai/trainer/train_callback.py

@@ -79,30 +79,11 @@ class GradientClippingCallback(TrainCallback):
     def __init__(self, max_grad_norm: float):
         self.max_grad_norm = max_grad_norm
 
-    def on_step_begin(self, context: TrainContext):
+    def on_step_end(self, context: TrainContext):
         _ = context
         clip_grad_norm_(context.model.parameters(), self.max_grad_norm)
 
 
-@CallbackFactory.register("scheduler")
-class SchedulerCallback(TrainCallback):
-    """
-    Scheduler callback for trainer.
-    """
-
-    def __init__(self):
-        pass
-
-    def on_train_begin(self, context: TrainContext):
-        for group in context.optimizer.param_groups:
-            if "initial_lr" not in group:
-                group["initial_lr"] = group["lr"]
-
-    def on_batch_end(self, context: TrainContext):
-        if context.scheduler:
-            context.scheduler.step()
-
-
 @CallbackFactory.register("checkpoint")
 class CheckpointCallback(TrainCallback):
     """

astrai/trainer/trainer.py

@@ -1,4 +1,5 @@
 import logging
+from itertools import batched
 from typing import List, Optional
 
 from astrai.config import TrainConfig
@@ -30,7 +31,6 @@ class Trainer:
             CallbackFactory.create("checkpoint", cfg.ckpt_dir, cfg.ckpt_interval),
             CallbackFactory.create("metric_logger", cfg.ckpt_dir, cfg.ckpt_interval),
             CallbackFactory.create("gradient_clipping", cfg.max_grad_norm),
-            CallbackFactory.create("scheduler"),
         ]
 
     def _build_context(self, checkpoint: Optional[Checkpoint]) -> TrainContext:
@@ -62,31 +62,32 @@ class Trainer:
 
         try:
             context.model.train()
-            # 1.epoch
+            accumulation_steps = max(self.train_config.accumulation_steps, 1)
+
             for epoch in range(context.epoch, self.train_config.n_epoch):
                 context.epoch = epoch
                 self._call_callbacks("on_epoch_begin", context)
 
-                accumulation_steps = max(self.train_config.accumulation_steps, 1)
+                for steps in batched(context.dataloader, accumulation_steps):
-                for batch in context.dataloader:
+                    self._call_callbacks("on_step_begin", context)
-                    if context.iteration % accumulation_steps == 0:
-                        # 2. step
-                        self._call_callbacks("on_step_begin", context)
-                        context.optimizer.step()
-                        context.optimizer.zero_grad()
-                        self._call_callbacks("on_step_end", context)
 
-                    # 3. batch
+                    step_batch_nums = len(steps)
-                    self._call_callbacks("on_batch_begin", context)
+                    for batch in steps:
-                    loss = context.strategy(batch)
+                        self._call_callbacks("on_batch_begin", context)
-                    context.loss = loss.item()
+                        loss = context.strategy(batch)
-                    context.iteration += 1
+                        context.loss = loss.item()
+                        context.iteration += 1
 
-                    # to make the loss normalized by accumulation steps
+                        stand_loss = loss / step_batch_nums
-                    stand_loss = loss / accumulation_steps
+                        stand_loss.backward()
-                    stand_loss.backward()
+                        self._call_callbacks("on_batch_end", context)
 
-                    self._call_callbacks("on_batch_end", context)
+                    self._call_callbacks("on_step_end", context)
+                    context.optimizer.step()
+                    context.optimizer.zero_grad()
+
+                    if context.scheduler:
+                        context.scheduler.step()
 
                 self._call_callbacks("on_epoch_end", context)
 

scripts/demo/generate_batch.py

@@ -24,12 +24,23 @@ def batch_generate():
         "请问什么是显卡",
     ]
 
+    prompts = [
+        tokenizer.apply_chat_template(
+            [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": q},
+            ],
+            tokenize=False,
+        )
+        for q in inputs
+    ]
+
     engine = InferenceEngine(
         model=model,
         tokenizer=tokenizer,
     )
     responses = engine.generate(
-        prompt=inputs,
+        prompt=prompts,
         stream=False,
         max_tokens=2048,
         temperature=0.8,

scripts/tools/train.py

@@ -155,18 +155,20 @@ def parse_args() -> argparse.Namespace:
 
 def ddp_wrap(model: nn.Module):
     local_rank = get_rank()
-    model = model.to(dtype=torch.bfloat16)
     ddp_model = DDP(
         model,
         device_ids=[local_rank],
         output_device=local_rank,
+        static_graph=True,
         find_unused_parameters=False,
+        gradient_as_bucket_view=True,
+        broadcast_buffers=False,
     )
     return ddp_model
 
 
 def create_optimizer(model: nn.Module, **kwargs) -> optim.Optimizer:
-    return optim.AdamW(model.parameters(), **kwargs)
+    return optim.AdamW(model.parameters(), fused=True, **kwargs)
 
 
 def create_scheduler(
@@ -231,6 +233,8 @@ def train(
         state_dict = st.load_file(weights_path)
         model.load_state_dict(state_dict, strict=False)
 
+    model = model.to(dtype=torch.bfloat16)
+
     strategy_kwargs = {
         "dpo_beta": dpo_beta,
         "label_smoothing": label_smoothing,

tests/inference/test_sample.py

@@ -48,12 +48,12 @@ def test_top_k_skip_when_zero():
 
 
 def test_top_k_batch_tensor():
-    """When top_k is a batch tensor, max element governs k for all rows."""
+    """Each row respects its own top_k."""
     logits = torch.tensor([[0.1, 0.5, 0.3], [0.9, 0.2, 0.1]])
     s = TopKStrategy(top_k=torch.tensor([2, 1]))
     result = s.apply(logits.clone(), filter_value=-1e9)
     assert (result[0] > -1e9).sum() == 2
-    assert (result[1] > -1e9).sum() == 2
+    assert (result[1] > -1e9).sum() == 1
 
 
 def test_top_p_nucleus_filtering():