From 319f6d7b44ed3f043f747a4f8a91f0205a88935a Mon Sep 17 00:00:00 2001
From: Brad Hilton <brad.hilton.nw@gmail.com>
Date: Wed, 15 Jul 2026 01:21:33 +0000
Subject: [PATCH 1/3] Upgrade ty to 0.0.59

---
 pyproject.toml                                |  5 +-
 src/art/gather.py                             | 11 +--
 src/art/guided_completion.py                  | 16 +++-
 src/art/langgraph/message_utils.py            |  6 +-
 src/art/local/backend.py                      | 24 +++--
 src/art/loss.py                               |  2 +-
 src/art/mcp/generate_scenarios.py             | 14 ++-
 src/art/megatron/context_parallel/executor.py | 14 +--
 src/art/megatron/dsv4/bridge.py               |  6 +-
 .../dsv4/kernel/tilelang_indexer_bwd.py       |  6 +-
 .../dsv4/kernel/tilelang_indexer_fwd.py       |  6 +-
 .../dsv4/kernel/tilelang_sparse_mla_bwd.py    |  6 +-
 .../dsv4/kernel/tilelang_sparse_mla_fwd.py    |  6 +-
 src/art/megatron/dsv4/lora.py                 |  2 +-
 src/art/megatron/dsv4/tokenizer.py            |  2 +-
 .../megatron/flex_attn/flash_dlse_patch.py    |  2 +-
 src/art/megatron/gdn/operator.py              |  2 +
 src/art/megatron/gdn/segment_layout.py        |  2 +
 src/art/megatron/lora.py                      | 36 ++++---
 .../megatron/model_support/handlers/dsv4.py   |  4 +-
 .../megatron/model_support/handlers/gemma4.py |  4 +-
 .../model_support/handlers/gpt_oss.py         |  2 +-
 .../model_support/handlers/qwen3_5.py         |  4 +-
 .../model_support/handlers/qwen3_common.py    |  2 +-
 src/art/megatron/prefix_tree_state.py         | 51 +++++-----
 src/art/megatron/runtime/bridge_runtime.py    | 10 +-
 src/art/megatron/train.py                     | 16 +++-
 src/art/megatron/training/microbatches.py     |  6 +-
 .../megatron/weights/merged_weight_export.py  |  2 +-
 src/art/pipeline_trainer/trainer.py           | 26 +++--
 src/art/pipeline_trainer/types.py             |  6 +-
 src/art/serverless/backend.py                 |  4 +-
 src/art/tinker_native/backend.py              |  5 +-
 src/art/trainer_rank/__init__.py              |  2 +-
 src/art/trajectories.py                       |  4 +-
 src/art/unsloth/train.py                      |  6 +-
 src/art/utils/format_message.py               | 23 +++--
 src/art/utils/lifecycle.py                    | 16 +++-
 .../calculate_step_metrics.py                 |  2 +-
 src/art/utils/trajectory_migration.py         |  6 +-
 tests/integration.py                          |  3 +-
 .../megatron_attention_oracle_worker.py       |  2 +-
 .../test_real_gdn_native_fla_cp.py            |  5 +-
 .../megatron/lora/test_dynamic_lora_slots.py  |  8 +-
 .../model_support/chat_template_rollout.py    |  5 +-
 .../megatron/model_support/oracle_harness.py  |  2 +-
 .../megatron/model_support/oracle_worker.py   | 10 +-
 .../model_support/routing_replay_bundle.py    |  8 +-
 .../megatron/train_inf_mismatch/real_path.py  |  5 +-
 .../trainability/yes_no_trainability.py       |  2 +-
 .../test_sft_last_assistant_tokenization.py   |  3 +
 tests/unit/test_dedicated_config.py           |  2 +-
 tests/unit/test_frontend_logging.py           |  2 +-
 tests/unit/test_moe_routing_real_path.py      |  2 +-
 tests/unit/test_moe_routing_replay.py         |  2 +-
 tests/unit/test_pipeline_trainer_batching.py  |  6 +-
 .../test_pipeline_trainer_local_backend.py    | 22 +++--
 tests/unit/test_pipeline_trainer_metrics.py   |  6 +-
 tests/unit/test_preprocessing_tokenize.py     | 16 ++--
 ...test_serverless_pipeline_trainer_compat.py |  8 +-
 tests/unit/test_tinker_renderers.py           |  2 +-
 .../test_tokenize_trajectory_groups.ipynb     |  4 +-
 tests/unit/test_track_api_cost.py             |  2 +-
 tests/unit/test_trainer_rank_validation.py    | 94 +++++++++----------
 tests/unit/test_trainer_rank_weird_shapes.py  | 36 +++----
 uv.lock                                       | 43 ++++-----
 66 files changed, 380 insertions(+), 289 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 8a0a0ae3e..3fa87c527 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -254,6 +254,7 @@ python-version = "3.12"
 # optional deps are installed. The ty:ignore comments are needed in CI (with all deps)
 # but become unused locally (without all deps).
 unused-ignore-comment = "ignore"
+unused-type-ignore-comment = "ignore"
 
 [tool.ty.analysis]
 # Allow unresolved imports for optional dependencies that may not be installed locally.
@@ -295,9 +296,11 @@ allowed-unresolved-imports = [
     "seaborn.**",
     # megatron deps
     "causal_conv1d.**",
+    "einops.**",
     "fla.**",
     "megatron.**",
     "quack.**",
+    "safetensors.**",
     "transformer_engine.**",
     "triton.**",
 ]
@@ -312,7 +315,7 @@ dev = [
     "pytest>=8.4.1",
     "nbval>=0.11.0",
     "pytest-xdist>=3.8.0",
-    "ty==0.0.14",
+    "ty==0.0.59",
     "pytest-asyncio>=1.1.0",
     "duckdb>=1.0.0",
     "pyarrow>=15.0.0",
diff --git a/src/art/gather.py b/src/art/gather.py
index 3a72df637..0dc53065a 100644
--- a/src/art/gather.py
+++ b/src/art/gather.py
@@ -3,7 +3,7 @@
 import contextlib
 import contextvars
 from dataclasses import dataclass, field
-from typing import Awaitable, Callable, Iterable, Iterator, Literal, overload
+from typing import Awaitable, Callable, Iterable, Iterator, Literal, Sequence, overload
 
 from openai.types.chat.chat_completion import Choice
 from tqdm import auto as tqdm
@@ -53,7 +53,7 @@ async def group_forward(g: Awaitable[TrajectoryGroup]):
         context.pbar.close()
 
     # Filter out any None results that may have been returned due to handled exceptions
-    processed_groups = []
+    processed_groups: list[TrajectoryGroup] = []
     for g in result_groups:
         if g is None:
             continue
@@ -113,12 +113,7 @@ async def gather_trajectories(
     pbar_desc: str | None = "gather",
     pbar_total_completion_tokens: bool = False,
     max_exceptions: int | float = 0,
-) -> (
-    list[Trajectory]
-    | list[Trajectory | BaseException]
-    | list[list[Trajectory]]
-    | list[list[Trajectory] | BaseException]
-):
+) -> Sequence[Trajectory | list[Trajectory] | BaseException]:
     if pbar_total_completion_tokens:
         print(
             "pbar_total_completion_tokens is deprecated and will be removed in a future version."
diff --git a/src/art/guided_completion.py b/src/art/guided_completion.py
index 6e8305c7d..f652a1d3f 100644
--- a/src/art/guided_completion.py
+++ b/src/art/guided_completion.py
@@ -13,7 +13,9 @@
 from pydantic import create_model
 
 
-def freeze_tool_schema(tool: dict, fixed_args: dict) -> ChatCompletionToolParam:
+def freeze_tool_schema(
+    tool: ChatCompletionToolParam, fixed_args: dict[str, object]
+) -> ChatCompletionToolParam:
     """
     Return a clone of *tool* whose parameters schema permits *only* `fixed_args`.
     Each field is cast to typing.Literal[value] so Pydantic emits an
@@ -27,7 +29,7 @@ def freeze_tool_schema(tool: dict, fixed_args: dict) -> ChatCompletionToolParam:
 
     locked = deepcopy(tool)
     locked["function"]["parameters"] = FrozenModel.model_json_schema()
-    return locked  # type: ignore
+    return locked
 
 
 def get_guided_completion_params(
@@ -36,7 +38,7 @@ def get_guided_completion_params(
 ) -> Tuple[
     List[str] | None,
     ChatCompletionToolChoiceOptionParam | None,
-    ChatCompletionToolParam | None,
+    List[ChatCompletionToolParam] | None,
 ]:
     """
     Given a completion from a teacher model, returns chat completion params that can be used to guide a student model's response.
@@ -48,7 +50,9 @@ def get_guided_completion_params(
 
     Returns a tuple of (guided_choice, tool_choice, tool_params).
     """
-    guided_choice, tool_choice, tool_params = None, None, None
+    guided_choice: List[str] | None = None
+    tool_choice: ChatCompletionToolChoiceOptionParam | None = None
+    tool_params: List[ChatCompletionToolParam] | None = None
 
     if (
         completion.choices[0].message.tool_calls
@@ -73,5 +77,7 @@ def get_guided_completion_params(
         ]
     else:
         content = completion.choices[0].message.content
+        if content is None:
+            raise ValueError("Completion has no text or tool call")
         guided_choice = [content]
-    return (guided_choice, tool_choice, tool_params)  # type: ignore
+    return (guided_choice, tool_choice, tool_params)
diff --git a/src/art/langgraph/message_utils.py b/src/art/langgraph/message_utils.py
index 3f9243c33..0c67b7c8b 100644
--- a/src/art/langgraph/message_utils.py
+++ b/src/art/langgraph/message_utils.py
@@ -1,5 +1,5 @@
 import json
-from typing import List, Union
+from typing import List, Union, cast
 
 from langchain_core.messages import (
     AIMessage,
@@ -88,7 +88,7 @@ def langchain_msg_to_openai(msg: BaseMessage):
 
 
 def convert_langgraph_messages(messages: List[object]) -> MessagesAndChoices:
-    converted = []
+    converted: MessagesAndChoices = []
 
     for msg in messages:
         response_metadata = getattr(msg, "response_metadata")
@@ -116,7 +116,7 @@ def convert_langgraph_messages(messages: List[object]) -> MessagesAndChoices:
         elif isinstance(msg, BaseMessage):
             converted.append(langchain_msg_to_openai(msg))
         elif isinstance(msg, dict):
-            converted.append(msg)
+            converted.append(cast(Message, msg))
         else:
             raise TypeError(f"Unsupported message type: {type(msg)}")
 
diff --git a/src/art/local/backend.py b/src/art/local/backend.py
index 8ec2c770c..9c8a9842e 100644
--- a/src/art/local/backend.py
+++ b/src/art/local/backend.py
@@ -75,6 +75,7 @@
 )
 from ..trajectories import Trajectory, TrajectoryGroup
 from ..types import (
+    Choice,
     LocalTrainResult,
     Message,
     TrainConfig,
@@ -189,6 +190,13 @@ def _tokenizer_cache_key(
     return (base_model, hashlib.sha256(chat_template.encode("utf-8")).hexdigest())
 
 
+def _load_training_tokenizer(base_model: str) -> PreTrainedTokenizerBase:
+    return cast(
+        PreTrainedTokenizerBase,
+        AutoTokenizer.from_pretrained(base_model),
+    )
+
+
 class LocalBackend(Backend):
     def __init__(
         self,
@@ -578,7 +586,7 @@ def _get_packed_tensors(
         tokenizer_key = _tokenizer_cache_key(model.base_model, internal_config)
         if tokenizer_key not in self._tokenizers:
             tokenizer = self._configure_training_tokenizer(
-                AutoTokenizer.from_pretrained(model.base_model),
+                _load_training_tokenizer(model.base_model),
                 model=model,
                 internal_config=internal_config,
             )
@@ -777,10 +785,10 @@ def _trajectory_log(self, trajectory: Trajectory) -> str:
         header = f"reward: {trajectory.reward} {' '.join(f'{k}: {v}' for k, v in trajectory.metrics.items())}\n\n"
         formatted_messages = []
         for message_or_choice in trajectory.messages_and_choices:
-            if isinstance(message_or_choice, dict):
-                message = message_or_choice
+            if isinstance(message_or_choice, Choice):
+                message = cast(Message, message_or_choice.message.model_dump())
             else:
-                message = cast(Message, message_or_choice.message.model_dump())  # ty:ignore[possibly-missing-attribute]
+                message = message_or_choice
             formatted_messages.append(format_message(message))
         return header + "\n".join(formatted_messages)
 
@@ -826,7 +834,7 @@ async def train(  # type: ignore[override]
         save_checkpoint: bool = True,
         # Verbosity
         verbose: bool = False,
-    ) -> LocalTrainResult:
+    ) -> LocalTrainResult:  # ty:ignore[invalid-method-override]
         """Train the model on the given trajectory groups.
 
         This method does NOT automatically log trajectories or metrics. Call
@@ -1071,7 +1079,7 @@ async def _train_model(
                     if hasattr(service, "register_lora_for_step"):
                         await service.register_lora_for_step(  # type: ignore[attr-defined]
                             next_step, next_checkpoint_dir
-                        )
+                        )  # ty:ignore[call-non-callable]
                     logger.info(
                         f"[BACKEND] _train_model SKIP: register_lora_for_step "
                         f"completed for step {next_step}"
@@ -1194,7 +1202,7 @@ async def _train_sft(
         tokenizer_key = _tokenizer_cache_key(model.base_model, internal_config)
         if tokenizer_key not in self._tokenizers:
             tokenizer = self._configure_training_tokenizer(
-                AutoTokenizer.from_pretrained(model.base_model),
+                _load_training_tokenizer(model.base_model),
                 model=model,
                 internal_config=internal_config,
             )
@@ -1364,7 +1372,7 @@ async def _experimental_pull_model_checkpoint(
                     f"No checkpoints found for {model.project}/{model.name} in local storage or S3"
                 )
             elif local_latest_step is None:
-                resolved_step = s3_latest_step  # type: ignore[assignment]
+                resolved_step = s3_latest_step  # type: ignore[assignment]  # ty:ignore[invalid-assignment]
                 if verbose:
                     print(f"Using latest checkpoint from S3: step {resolved_step}")
             elif s3_latest_step is None:
diff --git a/src/art/loss.py b/src/art/loss.py
index 99719be19..870b9377d 100644
--- a/src/art/loss.py
+++ b/src/art/loss.py
@@ -73,7 +73,7 @@ def align_inputs(self) -> AlignedLossInputs:
             weights=shift_tensor(inputs["weights"], 0.0),
             group_ids=shift_tensor(inputs["group_ids"], 0),
             original_logprobs=(
-                shift_tensor(inputs["original_logprobs"], 0.0)  # ty: ignore[invalid-key]
+                shift_tensor(inputs["original_logprobs"], 0.0)  # ty: ignore[invalid-key, invalid-argument-type]
                 if "original_logprobs" in inputs
                 else None
             ),
diff --git a/src/art/mcp/generate_scenarios.py b/src/art/mcp/generate_scenarios.py
index 6932d9c2e..ac349944d 100644
--- a/src/art/mcp/generate_scenarios.py
+++ b/src/art/mcp/generate_scenarios.py
@@ -5,6 +5,9 @@
 from typing import Any, Dict, List, Optional
 
 import openai
+from openai.types.shared_params.response_format_json_schema import (
+    ResponseFormatJSONSchema,
+)
 
 from art.mcp.types import GeneratedScenarioCollection, MCPResource, MCPTool
 from art.utils.logging import _C, dim, err, info, ok, step
@@ -137,7 +140,7 @@ async def generate_scenarios(
     if custom_instructions:
         prompt += f"\n\nPay close attention to the following instructions when generating scenarios:\n\n{custom_instructions}"
 
-    response_schema = {
+    response_schema: dict[str, object] = {
         "type": "object",
         "properties": {
             "scenarios": {
@@ -166,14 +169,15 @@ async def generate_scenarios(
     )
 
     t1 = time.perf_counter()
+    response_format: ResponseFormatJSONSchema = {
+        "type": "json_schema",
+        "json_schema": {"name": "scenario_list", "schema": response_schema},
+    }
     response = client_openai.chat.completions.create(
         model=generator_model,
         messages=[{"role": "user", "content": prompt}],
         max_completion_tokens=8000,
-        response_format={
-            "type": "json_schema",
-            "json_schema": {"name": "scenario_list", "schema": response_schema},
-        },
+        response_format=response_format,
     )
     dt = time.perf_counter() - t1
     ok(f"Model responded in {dt:.2f}s.")
diff --git a/src/art/megatron/context_parallel/executor.py b/src/art/megatron/context_parallel/executor.py
index 3636d7361..cc1de2a0b 100644
--- a/src/art/megatron/context_parallel/executor.py
+++ b/src/art/megatron/context_parallel/executor.py
@@ -2235,16 +2235,16 @@ def _run_context_parallel_backward(
             grad_outputs=tuple(stage_output_grads),
             allow_unused=True,
         )
-        grad_map = {
+        grad_map: dict[str, torch.Tensor | None] = {
             name: grad for name, grad in zip(input_names, input_grads, strict=True)
         }
         for grad_name in ("q_input", "k_input", "v_input"):
             grad_map[grad_name] = _sanitize_nested_stage_input_grad(
-                cast(torch.Tensor | None, grad_map.get(grad_name)),
+                grad_map.get(grad_name),
             )
         _scatter_stage_grad(
             target=dq_flat,
-            grad=cast(torch.Tensor | None, grad_map.get("q_input")),
+            grad=grad_map.get("q_input"),
             ranges=stage_plan.owner_local_q_ranges,
             state=state,
             head_major=True,
@@ -2252,14 +2252,14 @@ def _run_context_parallel_backward(
         if stage_plan.is_local_stage:
             _scatter_stage_grad(
                 target=dk_flat,
-                grad=cast(torch.Tensor | None, grad_map.get("k_input")),
+                grad=grad_map.get("k_input"),
                 ranges=stage_plan.owner_local_k_ranges,
                 state=state,
                 head_major=True,
             )
             _scatter_stage_grad(
                 target=dv_flat,
-                grad=cast(torch.Tensor | None, grad_map.get("v_input")),
+                grad=grad_map.get("v_input"),
                 ranges=stage_plan.owner_local_k_ranges,
                 state=state,
                 head_major=True,
@@ -2269,8 +2269,8 @@ def _run_context_parallel_backward(
             stage_record.clear()
             continue
         if not stage_plan.is_local_stage:
-            dk_remote = cast(torch.Tensor | None, grad_map.get("k_input"))
-            dv_remote = cast(torch.Tensor | None, grad_map.get("v_input"))
+            dk_remote = grad_map.get("k_input")
+            dv_remote = grad_map.get("v_input")
             if dk_remote is None:
                 dk_remote = k_flat.new_empty((k_flat.shape[0], 0, k_flat.shape[2]))
             if dv_remote is None:
diff --git a/src/art/megatron/dsv4/bridge.py b/src/art/megatron/dsv4/bridge.py
index d0eec5609..8cf02a910 100644
--- a/src/art/megatron/dsv4/bridge.py
+++ b/src/art/megatron/dsv4/bridge.py
@@ -669,7 +669,7 @@ def __init__(
 
         @property
         def group_key(self) -> str:
-            return cast(dict[str, str], self.export_hf_param)["gate"]
+            return self.export_hf_param["gate"]
 
         def hf_to_megatron(
             self,
@@ -718,7 +718,7 @@ def megatron_to_hf(self, megatron_weights: Any, megatron_module: Any):
             hf_param = cast(dict[str, str], self.hf_param)
             gate_suffix = hf_param["gate"].rpartition(".experts.")[2].split(".", 1)[1]
             remapped: dict[str, torch.Tensor] = {}
-            export = cast(dict[str, str], self.export_hf_param)
+            export = self.export_hf_param
             for gate_key, gate in converted.items():
                 if not gate_key.endswith(gate_suffix):
                     continue
@@ -1269,5 +1269,5 @@ def ensure_dsv4_bridge_registered() -> None:
         target=GPTModel,
         provider=MLAModelProvider,
         model_type="deepseek_v4",
-    )(ArtDeepSeekV4Bridge)
+    )(ArtDeepSeekV4Bridge)  # ty: ignore[invalid-argument-type]
     _DSV4_BRIDGE_REGISTERED = True
diff --git a/src/art/megatron/dsv4/kernel/tilelang_indexer_bwd.py b/src/art/megatron/dsv4/kernel/tilelang_indexer_bwd.py
index d17f17065..65cd35333 100644
--- a/src/art/megatron/dsv4/kernel/tilelang_indexer_bwd.py
+++ b/src/art/megatron/dsv4/kernel/tilelang_indexer_bwd.py
@@ -1,7 +1,5 @@
 # ruff: noqa
 # Adapted from miles_plugins/models/glm5/ops/tilelang_indexer_bwd.py for DeepSeek-V4.
-from typing import Any, cast
-
 import torch
 
 from art.megatron.dsv4.kernel.tilelang_import import (
@@ -12,8 +10,8 @@
 
 _tl, _T = import_tilelang()
 
-tl = cast(Any, _tl)
-T = cast(Any, _T)
+tl = _tl
+T = _T
 
 BF16 = T.bfloat16
 FP32 = T.float32
diff --git a/src/art/megatron/dsv4/kernel/tilelang_indexer_fwd.py b/src/art/megatron/dsv4/kernel/tilelang_indexer_fwd.py
index 1f8ff482a..63f3aec62 100644
--- a/src/art/megatron/dsv4/kernel/tilelang_indexer_fwd.py
+++ b/src/art/megatron/dsv4/kernel/tilelang_indexer_fwd.py
@@ -4,8 +4,6 @@
 #   - Operates on [seqlen, batch, heads, dim] (SBHD) layout, batch handled externally
 #   - Uses causal mask via cu_seqlens instead of variable-length packed sequences
 #   - Supports compressed KV (seq_len_kv = seq_len_q / compress_ratio)
-from typing import Any, cast
-
 import torch
 
 from art.megatron.dsv4.kernel.tilelang_import import (
@@ -16,8 +14,8 @@
 
 _tilelang, _T = import_tilelang()
 
-tilelang = cast(Any, _tilelang)
-T = cast(Any, _T)
+tilelang = _tilelang
+T = _T
 
 
 @tilelang.jit(
diff --git a/src/art/megatron/dsv4/kernel/tilelang_sparse_mla_bwd.py b/src/art/megatron/dsv4/kernel/tilelang_sparse_mla_bwd.py
index a9c2c743c..b854f5063 100644
--- a/src/art/megatron/dsv4/kernel/tilelang_sparse_mla_bwd.py
+++ b/src/art/megatron/dsv4/kernel/tilelang_sparse_mla_bwd.py
@@ -5,8 +5,6 @@
 #   - Single-head KV: kv shape [B, S_kv, D] (no kv_group, no D/D_tail split)
 #   - Index shape: [B, S, topk] (no kv_group dim)
 #   - Outputs: dQ [B, S, H, D], dKV [B, S_kv, D], dAttnSink [H]
-from typing import Any, cast
-
 import torch
 
 from art.megatron.dsv4.kernel.tilelang_import import (
@@ -17,8 +15,8 @@
 
 _tilelang, _T = import_tilelang()
 
-tilelang = cast(Any, _tilelang)
-T = cast(Any, _T)
+tilelang = _tilelang
+T = _T
 
 
 @tilelang.jit(out_idx=[-1])
diff --git a/src/art/megatron/dsv4/kernel/tilelang_sparse_mla_fwd.py b/src/art/megatron/dsv4/kernel/tilelang_sparse_mla_fwd.py
index 34cf531db..860b27c06 100644
--- a/src/art/megatron/dsv4/kernel/tilelang_sparse_mla_fwd.py
+++ b/src/art/megatron/dsv4/kernel/tilelang_sparse_mla_fwd.py
@@ -5,8 +5,6 @@
 #   - Single-head KV: kv shape [B, S_kv, D] (no kv_group, no D/D_tail split)
 #   - Index shape: [B, S, topk] (no kv_group dim)
 #   - Output: [B, S, H, D] + LSE [B, S, H]
-from typing import Any, cast
-
 import torch
 
 from art.megatron.dsv4.kernel.tilelang_import import (
@@ -17,8 +15,8 @@
 
 _tilelang, _T = import_tilelang()
 
-tilelang = cast(Any, _tilelang)
-T = cast(Any, _T)
+tilelang = _tilelang
+T = _T
 
 
 @tilelang.jit(
diff --git a/src/art/megatron/dsv4/lora.py b/src/art/megatron/dsv4/lora.py
index 658077081..7b52c38f6 100644
--- a/src/art/megatron/dsv4/lora.py
+++ b/src/art/megatron/dsv4/lora.py
@@ -619,7 +619,7 @@ def static_config_run(*args: Any, **kwargs: Any) -> Any:
             cache[key] = _dsv4_te_permutation_config(int(key[0]))
         return original_run(*args, **kwargs)
 
-    static_config_run._art_dsv4_static_config_wrapped = True  # type: ignore[attr-defined]
+    static_config_run._art_dsv4_static_config_wrapped = True  # type: ignore[attr-defined]  # ty:ignore[unresolved-attribute]
     autotuner.run = static_config_run
     autotuner._art_dsv4_static_config_wrapped = True
 
diff --git a/src/art/megatron/dsv4/tokenizer.py b/src/art/megatron/dsv4/tokenizer.py
index 1fb33adf2..e9eed3f90 100644
--- a/src/art/megatron/dsv4/tokenizer.py
+++ b/src/art/megatron/dsv4/tokenizer.py
@@ -28,7 +28,7 @@ def get_dsv4_tokenizer(
     added_vocab_size = len(added_vocab)
     tokenizer_vocab_size = tokenizer.vocab_size
 
-    class _ArtDsv4Tokenizer(tokenizer.__class__):  # type: ignore[misc, valid-type]
+    class _ArtDsv4Tokenizer(tokenizer.__class__):  # type: ignore[misc, valid-type]  # ty:ignore[unsupported-base]
         def apply_chat_template(
             self,
             messages: list[dict[str, Any]],
diff --git a/src/art/megatron/flex_attn/flash_dlse_patch.py b/src/art/megatron/flex_attn/flash_dlse_patch.py
index 560a14ef9..29eee3203 100644
--- a/src/art/megatron/flex_attn/flash_dlse_patch.py
+++ b/src/art/megatron/flex_attn/flash_dlse_patch.py
@@ -138,7 +138,7 @@ def bwd_postprocess_convert_art(
             cluster_size=cluster_size,
         )
 
-    bwd_postprocess_convert_art.compile_cache = (  # type: ignore[attr-defined]
+    bwd_postprocess_convert_art.compile_cache = (  # type: ignore[attr-defined]  # ty:ignore[unresolved-attribute]
         original_bwd_postprocess_convert.compile_cache
     )
     cute_interface_any._bwd_postprocess_convert = bwd_postprocess_convert_art
diff --git a/src/art/megatron/gdn/operator.py b/src/art/megatron/gdn/operator.py
index 302f43380..d593d0f12 100644
--- a/src/art/megatron/gdn/operator.py
+++ b/src/art/megatron/gdn/operator.py
@@ -1,3 +1,5 @@
+# ty: ignore[invalid-argument-type, unknown-argument]
+
 from __future__ import annotations
 
 from types import MethodType
diff --git a/src/art/megatron/gdn/segment_layout.py b/src/art/megatron/gdn/segment_layout.py
index 0b5feb760..c12140967 100644
--- a/src/art/megatron/gdn/segment_layout.py
+++ b/src/art/megatron/gdn/segment_layout.py
@@ -1,3 +1,5 @@
+# ty: ignore[invalid-argument-type, unknown-argument]
+
 from __future__ import annotations
 
 from typing import Any
diff --git a/src/art/megatron/lora.py b/src/art/megatron/lora.py
index e6b29cb97..295c2afc9 100644
--- a/src/art/megatron/lora.py
+++ b/src/art/megatron/lora.py
@@ -874,8 +874,8 @@ def _collect_templates(
                 for suffix, param in module._lora_params(slot_ref):
                     if not module._should_export_parameter(param):
                         continue
-                    sharded = bool(param.lora_tp_sharded)  # type: ignore[attr-defined]
-                    shard_domain = param.lora_shard_domain  # type: ignore[attr-defined]
+                    sharded = bool(param.lora_tp_sharded)  # type: ignore[attr-defined]  # ty:ignore[unresolved-attribute]
+                    shard_domain = param.lora_shard_domain  # type: ignore[attr-defined]  # ty:ignore[unresolved-attribute]
                     templates.append(
                         _LoraPublishTemplate(
                             adapter_model_prefix=module.adapter_model_prefix,
@@ -1056,7 +1056,9 @@ def _expert_grouped_lora_dual_forward(
     if isinstance(counts, list):
         counts = torch.tensor(counts, dtype=torch.int64, device="cpu")
     if x.shape[0] == 0:
-        return x.new_zeros((x.shape[0], module.linear_fc1.out_features))
+        return x.new_zeros(
+            (x.shape[0], module.linear_fc1.out_features)  # ty: ignore[unresolved-attribute]
+        )
     gate = module.gate_lora.active_lora_tensors()
     up = module.up_lora.active_lora_tensors()
     if gate is None or up is None:
@@ -1478,7 +1480,7 @@ def __init__(
             self.lora = _parallel_lora(
                 adapter_model_prefix=f"{adapter_model_prefix}.{{expert}}.gate_up_proj",
                 linear=linear_fc1,
-                out_features=linear_fc1.out_features,
+                out_features=linear_fc1.out_features,  # ty: ignore[unresolved-attribute]
                 rank=rank,
                 alpha=alpha,
                 layout="column",
@@ -1487,7 +1489,7 @@ def __init__(
                 allreduce=False,
                 num_local_experts=num_local_experts,
             )
-            gate_out_features = linear_fc1.out_features // 2
+            gate_out_features = linear_fc1.out_features // 2  # ty: ignore[unresolved-attribute]
             expert_tp_world_size = _get_shard_world_size("expert_tp")
             _set_lora_shard_strategy_metadata(
                 self.lora.B_T,
@@ -1501,7 +1503,7 @@ def __init__(
             self.gate_lora, self.up_lora = _parallel_lora_pair(
                 adapter_model_prefix=f"{adapter_model_prefix}.{{expert}}",
                 linear=linear_fc1,
-                out_features=linear_fc1.out_features // 2,
+                out_features=linear_fc1.out_features // 2,  # ty: ignore[unresolved-attribute]
                 rank=rank,
                 alpha=alpha,
                 layout="column",
@@ -1521,7 +1523,10 @@ def forward(
         )(x, tokens_per_expert)
         adapter_out = (
             _expert_grouped_lora_forward(
-                self.lora, x, tokens_per_expert, self.linear_fc1.out_features
+                self.lora,
+                x,
+                tokens_per_expert,
+                self.linear_fc1.out_features,  # ty: ignore[unresolved-attribute]
             )
             if self.fused_gate_up
             else _expert_grouped_lora_dual_forward(self, x, tokens_per_expert)
@@ -1543,7 +1548,7 @@ def __init__(
         self.lora = _parallel_lora(
             adapter_model_prefix=f"{adapter_model_prefix}.{{expert}}.down_proj",
             linear=linear_fc2,
-            out_features=linear_fc2.out_features,
+            out_features=linear_fc2.out_features,  # ty: ignore[unresolved-attribute]
             rank=rank,
             alpha=alpha,
             layout="row",
@@ -1564,7 +1569,10 @@ def forward(
             self.linear_fc2,
         )(x, tokens_per_expert)
         adapter_out = _expert_grouped_lora_forward(
-            self.lora, x, tokens_per_expert, self.linear_fc2.out_features
+            self.lora,
+            x,
+            tokens_per_expert,
+            self.linear_fc2.out_features,  # ty: ignore[unresolved-attribute]
         )
         # the reason there is no TP comm here is because the MoE token routing handles
         # expert TP comm externally
@@ -1703,7 +1711,7 @@ def wrap_standard_self_attention(
             "linear_qkv",
             TELayerNormColumnParallelLinear,
         )
-        self_attention.linear_qkv = SelfAttentionLinearQKVLoRA(
+        self_attention.linear_qkv = SelfAttentionLinearQKVLoRA(  # ty: ignore[invalid-assignment]
             adapter_model_prefix=f"{adapter_model_prefix}.self_attn",
             linear_qkv=self_attention_linear_qkv,
             rank=rank,
@@ -1768,9 +1776,9 @@ def wrap_grouped_moe_experts(
         mlp_experts_linear_fc1 = _unwrap_attr(
             experts.linear_fc1,
             "linear_fc1",
-            TEColumnParallelGroupedLinear,  # type: ignore[arg-type]
+            TEColumnParallelGroupedLinear,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
         )
-        experts.linear_fc1 = MLPExpertsLinearFC1LoRA(
+        experts.linear_fc1 = MLPExpertsLinearFC1LoRA(  # ty: ignore[invalid-assignment]
             adapter_model_prefix=f"{adapter_model_prefix}.mlp.experts",
             linear_fc1=mlp_experts_linear_fc1,
             rank=rank,
@@ -1785,9 +1793,9 @@ def wrap_grouped_moe_experts(
         linear_fc2 = _unwrap_attr(
             experts.linear_fc2,
             "linear_fc2",
-            TERowParallelGroupedLinear,  # type: ignore[arg-type]
+            TERowParallelGroupedLinear,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
         )
-        experts.linear_fc2 = MLPExpertsLinearFC2LoRA(
+        experts.linear_fc2 = MLPExpertsLinearFC2LoRA(  # ty: ignore[invalid-assignment]
             adapter_model_prefix=f"{adapter_model_prefix}.mlp.experts",
             linear_fc2=linear_fc2,
             rank=rank,
diff --git a/src/art/megatron/model_support/handlers/dsv4.py b/src/art/megatron/model_support/handlers/dsv4.py
index c6a604049..213683c15 100644
--- a/src/art/megatron/model_support/handlers/dsv4.py
+++ b/src/art/megatron/model_support/handlers/dsv4.py
@@ -268,10 +268,10 @@ def preprocess_hook(
                     )
                 return tuple(preproc_output)
 
-            gpt_module._preprocess = preprocess_hook  # type: ignore[attr-defined]
+            gpt_module._preprocess = preprocess_hook  # type: ignore[attr-defined]  # ty:ignore[invalid-assignment]
 
     def collect_layer_families(self, provider: Any) -> list[LayerFamilyInstance]:
-        ratios = list(getattr(provider, "dsv4_compress_ratios", ()) or ())
+        ratios: list[int] = list(getattr(provider, "dsv4_compress_ratios", ()) or ())
 
         def first_layer_index(ratio: int) -> int | None:
             try:
diff --git a/src/art/megatron/model_support/handlers/gemma4.py b/src/art/megatron/model_support/handlers/gemma4.py
index 6534538f3..65a09cbbe 100644
--- a/src/art/megatron/model_support/handlers/gemma4.py
+++ b/src/art/megatron/model_support/handlers/gemma4.py
@@ -799,7 +799,7 @@ def preprocess_hook(
             )
             return tuple(preproc_output)
 
-        gpt_module._preprocess = preprocess_hook  # type: ignore[attr-defined]
+        gpt_module._preprocess = preprocess_hook  # type: ignore[attr-defined]  # ty:ignore[invalid-assignment]
 
 
 def _install_gemma4_full_recompute_patch(model_chunks: Sequence[Any]) -> None:
@@ -1945,7 +1945,7 @@ def ensure_gemma4_text_only_bridge_registered() -> None:
     from megatron.bridge.models.gemma_vl.gemma4_vl_bridge import Gemma4VLBridge
     from megatron.core.models.gpt.gpt_model import GPTModel
 
-    @MegatronModelBridge.register_bridge(
+    @MegatronModelBridge.register_bridge(  # ty: ignore[invalid-argument-type]
         source="Gemma4ForConditionalGeneration",
         target=GPTModel,
         provider=Gemma4ModelProvider,
diff --git a/src/art/megatron/model_support/handlers/gpt_oss.py b/src/art/megatron/model_support/handlers/gpt_oss.py
index 57a6445a1..c4f8e5b6a 100644
--- a/src/art/megatron/model_support/handlers/gpt_oss.py
+++ b/src/art/megatron/model_support/handlers/gpt_oss.py
@@ -480,7 +480,7 @@ def preprocess_hook(
             )
             return tuple(preproc_output)
 
-        gpt_module._preprocess = preprocess_hook  # type: ignore[attr-defined]
+        gpt_module._preprocess = preprocess_hook  # type: ignore[attr-defined]  # ty:ignore[invalid-assignment]
 
 
 def _install_weighted_bias_quick_geglu_patch() -> None:
diff --git a/src/art/megatron/model_support/handlers/qwen3_5.py b/src/art/megatron/model_support/handlers/qwen3_5.py
index 9cf2c7064..7829be1e8 100644
--- a/src/art/megatron/model_support/handlers/qwen3_5.py
+++ b/src/art/megatron/model_support/handlers/qwen3_5.py
@@ -1195,7 +1195,7 @@ def ensure_qwen35_text_only_bridge_registered() -> None:
     )
     from megatron.core.models.gpt.gpt_model import GPTModel
 
-    @MegatronModelBridge.register_bridge(
+    @MegatronModelBridge.register_bridge(  # ty: ignore[invalid-argument-type]
         source=_QWEN3_5_DENSE_HF_CLASS_NAME,
         target=GPTModel,
         provider=Qwen35VLModelProvider,
@@ -1205,7 +1205,7 @@ class _ArtQwen35DenseTextOnlyBridge(Qwen35VLBridge):
         def mapping_registry(self) -> Any:
             return _qwen35_text_only_mapping_registry(Qwen35VLBridge)
 
-    @MegatronModelBridge.register_bridge(
+    @MegatronModelBridge.register_bridge(  # ty: ignore[invalid-argument-type]
         source=_QWEN3_5_MOE_HF_CLASS_NAME,
         target=GPTModel,
         provider=Qwen35VLMoEModelProvider,
diff --git a/src/art/megatron/model_support/handlers/qwen3_common.py b/src/art/megatron/model_support/handlers/qwen3_common.py
index 02497106d..74ba22b51 100644
--- a/src/art/megatron/model_support/handlers/qwen3_common.py
+++ b/src/art/megatron/model_support/handlers/qwen3_common.py
@@ -140,4 +140,4 @@ def preprocess_hook(*args, _preprocess=preprocess, **kwargs):
             )
             return tuple(preproc_output)
 
-        gpt_module._preprocess = preprocess_hook  # type: ignore[attr-defined]
+        gpt_module._preprocess = preprocess_hook  # type: ignore[attr-defined]  # ty:ignore[invalid-assignment]
diff --git a/src/art/megatron/prefix_tree_state.py b/src/art/megatron/prefix_tree_state.py
index cb8c2af3b..ef25b8897 100644
--- a/src/art/megatron/prefix_tree_state.py
+++ b/src/art/megatron/prefix_tree_state.py
@@ -208,34 +208,35 @@ def _build_sparse_prefix_tree_block_mask(
                 _empty_block_mask(seq_len=seq_len, block_size=block_size, device=device)
             )
             continue
-        row_masks.append(
-            build_block_mask_from_context(
-                FlexMaskSpec(
-                    q_len=seq_len,
-                    k_len=seq_len,
-                    block_size=block_size,
-                    slices=slices,
-                    exact_mask=ExactMaskMetadata(
-                        q_token_indices=token_indices,
-                        k_token_indices=token_indices,
-                        cache_key=(
-                            f"identity:{seq_len}"
-                            if sliding_window is None
-                            else f"identity:{seq_len}:sliding:{int(sliding_window)}"
-                        ),
+        row_mask = build_block_mask_from_context(
+            FlexMaskSpec(
+                q_len=seq_len,
+                k_len=seq_len,
+                block_size=block_size,
+                slices=slices,
+                exact_mask=ExactMaskMetadata(
+                    q_token_indices=token_indices,
+                    k_token_indices=token_indices,
+                    cache_key=(
+                        f"identity:{seq_len}"
+                        if sliding_window is None
+                        else f"identity:{seq_len}:sliding:{int(sliding_window)}"
                     ),
                 ),
-                context=prepare_block_mask_context(
-                    group_ids=group_ids_cpu[row_index],
-                    parent_ids=parent_ids_cpu[row_index],
-                    input_pos=None
-                    if input_pos_cpu is None
-                    else input_pos_cpu[row_index],
-                ),
-                sliding_window=sliding_window,
-                device=device,
-            )
+            ),
+            context=prepare_block_mask_context(
+                group_ids=group_ids_cpu[row_index],
+                parent_ids=parent_ids_cpu[row_index],
+                input_pos=None if input_pos_cpu is None else input_pos_cpu[row_index],
+            ),
+            sliding_window=sliding_window,
+            device=device,
         )
+        if row_mask is None:
+            row_mask = _empty_block_mask(
+                seq_len=seq_len, block_size=block_size, device=device
+            )
+        row_masks.append(row_mask)
     if not row_masks:
         return _empty_block_mask(seq_len=seq_len, block_size=block_size, device=device)
     return _stack_row_block_masks(
diff --git a/src/art/megatron/runtime/bridge_runtime.py b/src/art/megatron/runtime/bridge_runtime.py
index 8aabed47b..269087b96 100644
--- a/src/art/megatron/runtime/bridge_runtime.py
+++ b/src/art/megatron/runtime/bridge_runtime.py
@@ -657,9 +657,15 @@ def _replicated_hf_to_megatron(
     ):
         broadcast_device = _materialization_device()
     if self.tp_rank == 0:
-        tensor = hf_weights.to(device=cast(Any, broadcast_device), non_blocking=True)
+        tensor = hf_weights.to(
+            device=cast(Any, broadcast_device),  # ty: ignore[redundant-cast]
+            non_blocking=True,
+        )
     else:
-        tensor = torch.empty_like(hf_weights, device=cast(Any, broadcast_device))
+        tensor = torch.empty_like(
+            hf_weights,
+            device=cast(Any, broadcast_device),  # ty: ignore[redundant-cast]
+        )
     return self.broadcast_tensor_to_tp_ranks(tensor, src_rank=0)
 
 
diff --git a/src/art/megatron/train.py b/src/art/megatron/train.py
index a1130fde2..7de21a9c9 100644
--- a/src/art/megatron/train.py
+++ b/src/art/megatron/train.py
@@ -920,7 +920,7 @@ def load_adapter_into_model(
         for chunk in model_chunks:
             for module in chunk.modules():
                 if hasattr(module, "load_lora"):
-                    module.load_lora(adapter_model)  # type: ignore[attr-defined]
+                    module.load_lora(adapter_model)  # type: ignore[attr-defined]  # ty:ignore[call-non-callable]
 
     if optimizer is None:
         return
@@ -1281,13 +1281,16 @@ def run_megatron_sft_step(
     if not micro_inputs:
         raise ValueError("run_megatron_sft_step requires at least one trajectory")
 
+    micro_sample_indices: list[int | None]
     if isinstance(sample_index, list):
         if len(sample_index) != len(micro_inputs):
             raise ValueError(
                 "sample_index list length must match number of micro inputs: "
                 f"{len(sample_index)} != {len(micro_inputs)}"
             )
-        micro_sample_indices: list[int | None] = sample_index
+        micro_sample_indices = [
+            int(index) if index is not None else None for index in sample_index
+        ]
     else:
         assert len(micro_inputs) == 1
         micro_sample_indices = [sample_index]
@@ -1312,7 +1315,7 @@ def run_megatron_sft_step(
     )
 
     for chunk in model_chunks:
-        chunk.zero_grad_buffer()  # type: ignore[call-non-callable]
+        chunk.zero_grad_buffer()  # type: ignore[call-non-callable]  # ty:ignore[call-non-callable]
 
     raw_loss_sum: torch.Tensor | None = None
     loss_inputs_for_count: list[dict[str, torch.Tensor] | PreparedSFTMicroInputs] = []
@@ -1426,13 +1429,16 @@ def run_training_step(
     if not micro_inputs:
         raise ValueError("run_training_step requires at least one packed sequence")
 
+    micro_sample_indices: list[int | None]
     if isinstance(sample_index, list):
         if len(sample_index) != len(micro_inputs):
             raise ValueError(
                 "sample_index list length must match number of micro inputs: "
                 f"{len(sample_index)} != {len(micro_inputs)}"
             )
-        micro_sample_indices: list[int | None] = sample_index
+        micro_sample_indices = [
+            int(index) if index is not None else None for index in sample_index
+        ]
     else:
         assert len(micro_inputs) == 1
         micro_sample_indices = [sample_index]
@@ -1464,7 +1470,7 @@ def run_training_step(
         cp_lookahead_state.pending_prepared_micro = None
 
     for chunk in model_chunks:
-        chunk.zero_grad_buffer()  # type: ignore[call-non-callable]
+        chunk.zero_grad_buffer()  # type: ignore[call-non-callable]  # ty:ignore[call-non-callable]
 
     micro_count = len(micro_inputs)
     raw_loss_sum: torch.Tensor | None = None
diff --git a/src/art/megatron/training/microbatches.py b/src/art/megatron/training/microbatches.py
index c3b515c26..a80d5529f 100644
--- a/src/art/megatron/training/microbatches.py
+++ b/src/art/megatron/training/microbatches.py
@@ -65,7 +65,7 @@ def selected_tensor(value: torch.Tensor) -> torch.Tensor:
             return selected.clone()
         return selected
 
-    return PackedTensors(  # type: ignore[call-arg]
+    return PackedTensors(  # type: ignore[call-arg]  # ty:ignore[missing-typed-dict-key]
         **{
             key: selected_tensor(value)
             for key, value in packed_tensors.items()
@@ -78,7 +78,7 @@ def selected_tensor(value: torch.Tensor) -> torch.Tensor:
 
 @torch.no_grad()
 def _clone_packed_tensors(inputs: PackedTensors) -> PackedTensors:
-    return PackedTensors(  # type: ignore[call-arg]
+    return PackedTensors(  # type: ignore[call-arg]  # ty:ignore[missing-typed-dict-key]
         **{
             key: value.clone()
             for key, value in inputs.items()
@@ -227,7 +227,7 @@ def _select_next_step_first_micro(
 def _move_inputs_to_device(inputs: PackedTensors, device: torch.device) -> None:
     for key, value in inputs.items():
         if isinstance(value, torch.Tensor):
-            inputs[key] = value.to(device)  # type: ignore[index]
+            inputs[key] = value.to(device)  # type: ignore[index]  # ty:ignore[invalid-key]
 
 
 def _count_trainable_tokens(inputs: LossInputs | DispatchedPackedTensors) -> float:
diff --git a/src/art/megatron/weights/merged_weight_export.py b/src/art/megatron/weights/merged_weight_export.py
index 4000bcf63..0ff3d5abd 100644
--- a/src/art/megatron/weights/merged_weight_export.py
+++ b/src/art/megatron/weights/merged_weight_export.py
@@ -300,7 +300,7 @@ def ensure_merged_weight_transfer_group(
 
     error: BaseException | None = None
     if _is_sender_rank(rank):
-        init_kwargs = {
+        init_kwargs: dict[str, object] = {
             "master_address": spec.init_info.master_address,
             "master_port": spec.init_info.master_port,
             "world_size": spec.init_info.world_size,
diff --git a/src/art/pipeline_trainer/trainer.py b/src/art/pipeline_trainer/trainer.py
index fa9cae017..8a9d214b8 100644
--- a/src/art/pipeline_trainer/trainer.py
+++ b/src/art/pipeline_trainer/trainer.py
@@ -9,7 +9,16 @@
 from pathlib import Path
 import signal
 import time
-from typing import Any, AsyncIterator, Generic, Iterable, TypeVar, cast
+from typing import (
+    Any,
+    AsyncIterator,
+    Generic,
+    Iterable,
+    Mapping,
+    Sequence,
+    TypeVar,
+    cast,
+)
 
 T = TypeVar("T")
 
@@ -36,7 +45,7 @@
 def _to_async_iterator(iterable: Iterable[T] | AsyncIterator[T]) -> AsyncIterator[T]:
     """Convert a sync Iterable to an AsyncIterator, or pass through if already async."""
     if isinstance(iterable, AsyncIterator):
-        return iterable
+        return cast(AsyncIterator[T], iterable)
 
     async def _iter():
         for item in iterable:
@@ -760,9 +769,12 @@ async def _run_eval(self, step: int) -> None:
             finally:
                 token.var.reset(token)
                 eval_elapsed = time.monotonic() - eval_started
-            splits: dict[str, list[art.Trajectory | art.TrajectoryGroup]]
-            if isinstance(result, dict):
-                splits = result
+            splits: Mapping[str, Sequence[art.Trajectory | art.TrajectoryGroup]]
+            if isinstance(result, Mapping):
+                splits = cast(
+                    Mapping[str, Sequence[art.Trajectory | art.TrajectoryGroup]],
+                    result,
+                )
             else:
                 splits = {"val": result}
 
@@ -809,7 +821,7 @@ async def _run_eval(self, step: int) -> None:
 
     @staticmethod
     def _normalize_eval_items(
-        items: list[art.Trajectory | art.TrajectoryGroup],
+        items: Sequence[art.Trajectory | art.TrajectoryGroup],
     ) -> tuple[list[TrajectoryGroup], list[art.Trajectory]]:
         if not items:
             return [], []
@@ -970,7 +982,7 @@ def _persist_state(self, training_step: int) -> None:
         self.model.merge_state({PIPELINE_STATE_KEY: payload})
 
     def _log_checkpoint_history(self, step: int, metrics: dict[str, float]) -> None:
-        row = {
+        row: dict[str, int | float | str] = {
             (key if key.startswith("checkpoint/") else f"checkpoint/{key}"): value
             for key, value in metrics.items()
             if value == value
diff --git a/src/art/pipeline_trainer/types.py b/src/art/pipeline_trainer/types.py
index 4b04891e2..9ef682648 100644
--- a/src/art/pipeline_trainer/types.py
+++ b/src/art/pipeline_trainer/types.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from collections.abc import Awaitable, Callable
+from collections.abc import Awaitable, Callable, Mapping, Sequence
 from typing import TypeVar
 
 import art
@@ -22,7 +22,7 @@
 EvalFn = Callable[
     [art.TrainableModel, int, ConfigT],
     Awaitable[
-        list[Trajectory | TrajectoryGroup]
-        | dict[str, list[Trajectory | TrajectoryGroup]]
+        Sequence[Trajectory | TrajectoryGroup]
+        | Mapping[str, Sequence[Trajectory | TrajectoryGroup]]
     ],
 ]
diff --git a/src/art/serverless/backend.py b/src/art/serverless/backend.py
index 892460395..ea83475d1 100644
--- a/src/art/serverless/backend.py
+++ b/src/art/serverless/backend.py
@@ -257,7 +257,7 @@ async def train(  # type: ignore[override]
         save_checkpoint: bool = True,
         # Verbosity
         verbose: bool = False,
-    ) -> ServerlessTrainResult:
+    ) -> ServerlessTrainResult:  # ty:ignore[invalid-method-override]
         """Train the model on the given trajectory groups.
 
         This method does NOT automatically log trajectories or metrics. Call
@@ -858,7 +858,7 @@ async def _experimental_push_to_s3(
 
             # Pull from W&B to local temp dir
             checkpoint_dir = await self._experimental_pull_model_checkpoint(
-                model,  # type: ignore[arg-type]
+                model,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
                 step=step,
                 verbose=verbose,
             )
diff --git a/src/art/tinker_native/backend.py b/src/art/tinker_native/backend.py
index 9728a290d..1701c6cc4 100644
--- a/src/art/tinker_native/backend.py
+++ b/src/art/tinker_native/backend.py
@@ -735,10 +735,7 @@ async def _build_model_state(self, model: TrainableModel) -> ModelState:
 
     def _resolve_model_config(self, model: TrainableModel) -> TinkerNativeModelConfig:
         internal_config = model._internal_config or {}
-        tinker_native_args = cast(
-            dev.TinkerNativeArgs | None,
-            internal_config.get("tinker_native_args"),
-        )
+        tinker_native_args = internal_config.get("tinker_native_args")
         renderer_name = (
             tinker_native_args.get("renderer_name")
             if tinker_native_args is not None
diff --git a/src/art/trainer_rank/__init__.py b/src/art/trainer_rank/__init__.py
index f7b51bd44..c60bb3a13 100644
--- a/src/art/trainer_rank/__init__.py
+++ b/src/art/trainer_rank/__init__.py
@@ -2876,7 +2876,7 @@ def _materialize(inputs: ForwardInputs) -> ForwardInputs:
 
 def _flatten(inputs: ForwardInputs) -> Iterator[AnyForwardInput]:
     if isinstance(inputs, ForwardInput):
-        yield inputs
+        yield inputs  # ty: ignore[invalid-yield]
         return
     for item in _nested_forward_children(inputs):
         yield from _flatten(item)
diff --git a/src/art/trajectories.py b/src/art/trajectories.py
index 4314a3733..b844ac41f 100644
--- a/src/art/trajectories.py
+++ b/src/art/trajectories.py
@@ -123,7 +123,7 @@ def get_messages(messages_and_choices: MessagesAndChoices) -> Messages:
             msg = dict(message_or_choice)
             if msg.get("content") is None:
                 msg["content"] = ""
-            messages.append(msg)  # type: ignore[arg-type]
+            messages.append(msg)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     return messages
 
 
@@ -223,7 +223,7 @@ def __deepcopy__(self, memo: dict[int, Any] | None = None):
     def log(self, message: str) -> None:
         self.logs.append(message)
 
-    def __iter__(self) -> Iterator[Trajectory]:  # type: ignore[override]
+    def __iter__(self) -> Iterator[Trajectory]:  # type: ignore[override]  # ty:ignore[invalid-method-override]
         return iter(self.trajectories)
 
     def __len__(self) -> int:
diff --git a/src/art/unsloth/train.py b/src/art/unsloth/train.py
index 3fe51823b..ca95d7d67 100644
--- a/src/art/unsloth/train.py
+++ b/src/art/unsloth/train.py
@@ -435,7 +435,7 @@ def compute_loss(
             f"Sequence length ({seq_len}) must be evenly divisible by chunk size ({chunk_size})"
         )
         os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1"
-        forward_kwargs = {}
+        forward_kwargs: dict[str, object] = {}
         if "pixel_values" in inputs:
             forward_kwargs["pixel_values"] = inputs["pixel_values"]
         if "image_grid_thw" in inputs:
@@ -602,7 +602,7 @@ def calculate_logprobs(
     trainer: "GRPOTrainer",
     input_ids: torch.Tensor,
     causal_mask: torch.Tensor,
-    forward_kwargs: dict[str, torch.Tensor],
+    forward_kwargs: dict[str, object],
     next_input_ids: torch.Tensor,
     lm_head_t: torch.Tensor,
     chunk_size: int,
@@ -830,7 +830,7 @@ async def run_unsloth_rl_training(
     for offset in range(0, packed_tensors["tokens"].shape[0]):
         for _ in range(2 if warmup else 1):
             if precalculate_logprobs and not warmup:
-                packed_tensors["original_logprobs"] = packed_tensors["logprobs"]  # type: ignore[index]
+                packed_tensors["original_logprobs"] = packed_tensors["logprobs"]  # type: ignore[index]  # ty:ignore[invalid-key]
                 packed_tensors["logprobs"] = _precalculate_new_logprobs(
                     ctx,
                     packed_tensors,
diff --git a/src/art/utils/format_message.py b/src/art/utils/format_message.py
index fd3fe20dc..2b06cdb98 100644
--- a/src/art/utils/format_message.py
+++ b/src/art/utils/format_message.py
@@ -1,23 +1,22 @@
-from openai.types.chat.chat_completion_message_function_tool_call import (
-    ChatCompletionMessageFunctionToolCall,
-)
+from openai.types.chat.chat_completion_message_param import ChatCompletionMessageParam
 
-from ..types import Message
 
-
-def format_message(message: Message) -> str:
+def format_message(message: ChatCompletionMessageParam) -> str:
     """Format a message into a readable string."""
     # Format the role and content
     role = message["role"].capitalize()
     content = message.get("content", message.get("refusal", "")) or ""
 
     # Format any tool calls
-    tool_calls_text = "\n" if content else ""
-    tool_calls_text += "\n".join(
-        f"{tool_call['function']['name']}({tool_call['function']['arguments']})"  # ty:ignore[invalid-key]
-        for tool_call in message.get("tool_calls") or []
-        if "function" in tool_call
-    )
+    tool_calls = []
+    for tool_call in message.get("tool_calls") or []:
+        if not isinstance(tool_call, dict):
+            continue
+        function = tool_call.get("function")
+        if not isinstance(function, dict):
+            continue
+        tool_calls.append(f"{function.get('name')}({function.get('arguments')})")
+    tool_calls_text = ("\n" if content else "") + "\n".join(tool_calls)
 
     # Combine all parts
     formatted_message = f"{role}:\n{content}{tool_calls_text}"
diff --git a/src/art/utils/lifecycle.py b/src/art/utils/lifecycle.py
index 6fe315659..fd171238d 100644
--- a/src/art/utils/lifecycle.py
+++ b/src/art/utils/lifecycle.py
@@ -9,6 +9,7 @@
 import subprocess
 import sys
 import time
+from types import FrameType
 from typing import Any
 
 
@@ -149,7 +150,9 @@ class ServiceLifecycle:
     def __init__(self) -> None:
         self.closing = False
         self._close_callback: Callable[[], None] | None = None
-        self._previous_signal_handlers: dict[int, Any] = {}
+        self._previous_signal_handlers: dict[
+            int, Callable[[int, FrameType | None], object] | int | None
+        ] = {}
 
     def begin_close(self) -> bool:
         if self.closing:
@@ -172,9 +175,16 @@ def _default_signal_exit(signum: int) -> None:
             previous = signal.getsignal(signum)
             self._previous_signal_handlers[signum] = previous
 
-            def _handler(received_signum, frame, *, _previous=previous):
+            def _handler(
+                received_signum: int,
+                frame: FrameType | None,
+                *,
+                _previous: Callable[[int, FrameType | None], object]
+                | int
+                | None = previous,
+            ) -> None:
                 close()
-                if callable(_previous):
+                if not isinstance(_previous, int) and _previous is not None:
                     _previous(received_signum, frame)
                     return
                 if _previous == signal.SIG_IGN:
diff --git a/src/art/utils/old_benchmarking/calculate_step_metrics.py b/src/art/utils/old_benchmarking/calculate_step_metrics.py
index be8b07f38..715f826e0 100644
--- a/src/art/utils/old_benchmarking/calculate_step_metrics.py
+++ b/src/art/utils/old_benchmarking/calculate_step_metrics.py
@@ -20,4 +20,4 @@ def calculate_step_std_dev(trajectory_groups: list[TrajectoryGroup]) -> float:
     if len(std_devs) == 0:
         return 0
 
-    return sum(std_devs) / len(std_devs)
+    return float(sum(std_devs) / len(std_devs))
diff --git a/src/art/utils/trajectory_migration.py b/src/art/utils/trajectory_migration.py
index 2cdf200cf..7ec947b69 100644
--- a/src/art/utils/trajectory_migration.py
+++ b/src/art/utils/trajectory_migration.py
@@ -79,9 +79,9 @@ def trajectory_to_dict(trajectory: Trajectory) -> dict[str, Any]:
 def message_or_choice_to_dict(message_or_choice: MessageOrChoice) -> dict[str, Any]:
     # messages are sometimes stored as dicts, so we need to handle both cases
     item_dict = (
-        message_or_choice
-        if isinstance(message_or_choice, dict)
-        else message_or_choice.to_dict()  # ty:ignore[possibly-missing-attribute]
+        message_or_choice.model_dump()
+        if isinstance(message_or_choice, Choice)
+        else message_or_choice
     )
 
     if "logprobs" in item_dict:
diff --git a/tests/integration.py b/tests/integration.py
index 3ed8070c5..7308a3649 100644
--- a/tests/integration.py
+++ b/tests/integration.py
@@ -322,8 +322,9 @@ def main() -> int:
         print("Running all notebooks")
 
     # Process notebook configurations
-    processed_configs = []
+    processed_configs: list[dict[str, Any]] = []
     for config in selected_configs:
+        processed_config: dict[str, Any]
         if isinstance(config, str):
             # Handle legacy string format
             processed_config = {"path": config, "variables": {}}
diff --git a/tests/integration/megatron/cp_attn/megatron_attention_oracle_worker.py b/tests/integration/megatron/cp_attn/megatron_attention_oracle_worker.py
index 697dfce64..0bdca2e15 100644
--- a/tests/integration/megatron/cp_attn/megatron_attention_oracle_worker.py
+++ b/tests/integration/megatron/cp_attn/megatron_attention_oracle_worker.py
@@ -97,7 +97,7 @@ def run_worker_subprocess(
         log_file.flush()
         live_log_file.write(header)
         live_log_file.flush()
-        run = subprocess.Popen(
+        run = subprocess.Popen(  # ty: ignore[no-matching-overload]
             command,
             cwd=str(worker_cwd),
             env=env,
diff --git a/tests/integration/megatron/gdn_shared_prefix/test_real_gdn_native_fla_cp.py b/tests/integration/megatron/gdn_shared_prefix/test_real_gdn_native_fla_cp.py
index 63c3ce43e..3d11115f0 100644
--- a/tests/integration/megatron/gdn_shared_prefix/test_real_gdn_native_fla_cp.py
+++ b/tests/integration/megatron/gdn_shared_prefix/test_real_gdn_native_fla_cp.py
@@ -239,11 +239,12 @@ def _native_gdn_cp_prepared_varlen_worker(
         ).requires_grad_(True)
         conv_grad = torch.randn_like(conv0_ref)
         rec_grad = torch.randn_like(rec0_ref)
+        assert cp_gdn.num_value_heads is not None
         output_grad = torch.randn(
             1,
             int(lengths.sum().item()),
-            cast(int, cp_gdn.num_value_heads) // cast(int, cp_gdn.tp_size),
-            cast(int, cp_gdn.value_head_dim),
+            cp_gdn.num_value_heads // cp_gdn.tp_size,
+            cp_gdn.value_head_dim,
             device=hidden.device,
             dtype=GDN_CORRECTNESS_DTYPE,
         )
diff --git a/tests/integration/megatron/lora/test_dynamic_lora_slots.py b/tests/integration/megatron/lora/test_dynamic_lora_slots.py
index ca7cf921b..c8d3335f5 100644
--- a/tests/integration/megatron/lora/test_dynamic_lora_slots.py
+++ b/tests/integration/megatron/lora/test_dynamic_lora_slots.py
@@ -62,7 +62,7 @@ def test_dynamic_lora_slots_capture_recompute_context_and_step_independently() -
         assert torch.allclose(actual, expected, atol=0, rtol=0)
         assert slot_a.rank == 1
         assert slot_a.scale == 32.0
-        assert lora._slot(ref_b).scale == 8.0  # type: ignore[union-attr]
+        assert lora._slot(ref_b).scale == 8.0  # type: ignore[union-attr]  # ty:ignore[unresolved-attribute]
 
         trainer = _trainer_for(lora, device)
         cpu_adapter = {
@@ -326,8 +326,8 @@ def _assert_checkpoint_recomputes_with(
         y = checkpoint(lambda t: lora(t), *checkpoint_args, x)
     with use_lora_slot(ambient_ref):
         y.sum().backward()
-    assert lora._slot(expected_ref).A_T.grad is not None  # type: ignore[union-attr]
-    assert lora._slot(ambient_ref).A_T.grad is None  # type: ignore[union-attr]
+    assert lora._slot(expected_ref).A_T.grad is not None  # type: ignore[union-attr]  # ty:ignore[unresolved-attribute]
+    assert lora._slot(ambient_ref).A_T.grad is None  # type: ignore[union-attr]  # ty:ignore[unresolved-attribute]
 
 
 def _assert_step_updates_only(
@@ -375,7 +375,7 @@ def _assert_reload_replaces_slot_optimizer(
     assert ref.name not in trainer._dynamic_optimizers
     assert [tuple(param.shape) for param in new_params] == [(4, 3), (3, 5)]
     assert all(old is not new for old, new in zip(old_params, new_params, strict=True))
-    assert lora._slot(ref).rank == 3  # type: ignore[union-attr]
+    assert lora._slot(ref).rank == 3  # type: ignore[union-attr]  # ty:ignore[unresolved-attribute]
 
 
 def _trainer_for(lora: LoRA, device: torch.device) -> TrainerRank:
diff --git a/tests/integration/megatron/model_support/chat_template_rollout.py b/tests/integration/megatron/model_support/chat_template_rollout.py
index beb950513..9892c4cbb 100644
--- a/tests/integration/megatron/model_support/chat_template_rollout.py
+++ b/tests/integration/megatron/model_support/chat_template_rollout.py
@@ -109,9 +109,12 @@ def run_chat_template_rollout(base_model: str) -> ChatTemplateRolloutReport:
     tokenizer = backend._tokenizers.get(tokenizer_key)
     if tokenizer is None:
         from transformers import AutoTokenizer
+        from transformers.tokenization_utils_base import PreTrainedTokenizerBase
 
+        loaded_tokenizer = AutoTokenizer.from_pretrained(base_model)
+        assert isinstance(loaded_tokenizer, PreTrainedTokenizerBase)
         tokenizer = backend._configure_training_tokenizer(
-            AutoTokenizer.from_pretrained(base_model),
+            loaded_tokenizer,
             model=model,
             internal_config=internal_config,
         )
diff --git a/tests/integration/megatron/model_support/oracle_harness.py b/tests/integration/megatron/model_support/oracle_harness.py
index db2ad58a3..9b98aee54 100644
--- a/tests/integration/megatron/model_support/oracle_harness.py
+++ b/tests/integration/megatron/model_support/oracle_harness.py
@@ -152,7 +152,7 @@ def selected_oracle_objectives() -> list[OracleObjective]:
     if normalized == "all":
         return list(SUPPORTED_ORACLE_OBJECTIVES)
     if normalized in SUPPORTED_ORACLE_OBJECTIVES:
-        return [cast(OracleObjective, normalized)]
+        return [normalized]
     supported = ", ".join((*SUPPORTED_ORACLE_OBJECTIVES, "all"))
     raise ValueError(
         f"Unsupported {ORACLE_OBJECTIVE_ENV} value '{raw}'. "
diff --git a/tests/integration/megatron/model_support/oracle_worker.py b/tests/integration/megatron/model_support/oracle_worker.py
index 85fc0f596..a845bd370 100644
--- a/tests/integration/megatron/model_support/oracle_worker.py
+++ b/tests/integration/megatron/model_support/oracle_worker.py
@@ -525,7 +525,7 @@ def _apply_requested_flex_backend_patch(flex_backend: str | None):
     else:
         raise RuntimeError(f"Unsupported flex backend request: {flex_backend}")
 
-    compiled_flex_attention._FORCED_FLEX_BACKEND = patched_backend  # type: ignore[invalid-assignment]
+    compiled_flex_attention._FORCED_FLEX_BACKEND = patched_backend  # type: ignore[invalid-assignment]  # ty:ignore[invalid-assignment]
     compiled_flex_attention._FORCED_FLEX_KERNEL_OPTIONS = patched_kernel_options
     compiled_flex_attention.dense_compiled_flex_attention = torch.compile(
         compiled_flex_attention._forced_flex_attention_dense
@@ -699,7 +699,7 @@ def _assert_runtime_configuration(
     try:
         from megatron.core.ssm.gated_delta_net import GatedDeltaNet
     except ImportError:  # pragma: no cover - optional dependency guard.
-        GatedDeltaNet = ()  # type: ignore[assignment]
+        GatedDeltaNet = ()  # type: ignore[assignment]  # ty:ignore[invalid-assignment]
     from megatron.core.transformer.attention import SelfAttention
 
     for chunk in model_chunks:
@@ -1044,7 +1044,7 @@ def _mutated_sanitize(grad: torch.Tensor | None) -> torch.Tensor | None:
         view.copy_(grad)
         return view
 
-    executor._sanitize_nested_stage_input_grad = _mutated_sanitize  # type: ignore[invalid-assignment]
+    executor._sanitize_nested_stage_input_grad = _mutated_sanitize  # type: ignore[invalid-assignment]  # ty:ignore[invalid-assignment]
     try:
         yield
     finally:
@@ -1066,8 +1066,8 @@ def _apply_attention_lse_normalize_mutation(mutation: SensitivityMutation | None
     def _identity(lse: torch.Tensor, **_kwargs: Any) -> torch.Tensor:
         return lse
 
-    compiled_flex_attention.normalize_flex_lse = _identity  # type: ignore[invalid-assignment]
-    executor.normalize_flex_lse = _identity  # type: ignore[invalid-assignment]
+    compiled_flex_attention.normalize_flex_lse = _identity  # type: ignore[invalid-assignment]  # ty:ignore[invalid-assignment]
+    executor.normalize_flex_lse = _identity  # type: ignore[invalid-assignment]  # ty:ignore[invalid-assignment]
     try:
         yield
     finally:
diff --git a/tests/integration/megatron/model_support/routing_replay_bundle.py b/tests/integration/megatron/model_support/routing_replay_bundle.py
index 379ca36ba..3990b07f3 100644
--- a/tests/integration/megatron/model_support/routing_replay_bundle.py
+++ b/tests/integration/megatron/model_support/routing_replay_bundle.py
@@ -161,7 +161,13 @@ def _dedupe_checkpoint_router_calls(
             previous_call_key == call_key
             and previous_route is not None
             and torch.equal(compact_route.expert_indices, previous_route.expert_indices)
-            and torch.equal(compact_route.expert_probs, previous_route.expert_probs)
+            and (
+                compact_route.expert_probs is None
+                and previous_route.expert_probs is None
+                or compact_route.expert_probs is not None
+                and previous_route.expert_probs is not None
+                and torch.equal(compact_route.expert_probs, previous_route.expert_probs)
+            )
             and torch.equal(compact_route.expert_mask, previous_route.expert_mask)
         )
         if is_checkpoint_duplicate:
diff --git a/tests/integration/megatron/train_inf_mismatch/real_path.py b/tests/integration/megatron/train_inf_mismatch/real_path.py
index cd996cdb0..164f28fb6 100644
--- a/tests/integration/megatron/train_inf_mismatch/real_path.py
+++ b/tests/integration/megatron/train_inf_mismatch/real_path.py
@@ -363,6 +363,7 @@ async def _collect_real_trajectory_groups(
     config: RealPathConfig,
 ) -> list[Any]:
     from transformers import AutoTokenizer
+    from transformers.tokenization_utils_base import PreTrainedTokenizerBase
 
     import art
     from art.megatron.model_support.tokenizer import (
@@ -371,8 +372,10 @@ async def _collect_real_trajectory_groups(
 
     if config.rollouts_per_prompt < 2:
         raise ValueError("real-path mismatch requires at least two rollouts per prompt")
+    loaded_tokenizer = AutoTokenizer.from_pretrained(config.output_parity.base_model)
+    assert isinstance(loaded_tokenizer, PreTrainedTokenizerBase)
     tokenizer = configure_tokenizer_for_model_support(
-        AutoTokenizer.from_pretrained(config.output_parity.base_model),
+        loaded_tokenizer,
         base_model=config.output_parity.base_model,
         internal_config={
             "allow_unvalidated_arch": config.output_parity.allow_unvalidated_arch
diff --git a/tests/integration/megatron/trainability/yes_no_trainability.py b/tests/integration/megatron/trainability/yes_no_trainability.py
index f9a213148..55b516b66 100644
--- a/tests/integration/megatron/trainability/yes_no_trainability.py
+++ b/tests/integration/megatron/trainability/yes_no_trainability.py
@@ -118,7 +118,7 @@ def build_prompts() -> list[str]:
     prompt_count = _get_env_int("ART_MODEL_SUPPORT_YES_NO_PROMPT_COUNT", 8)
     if prompt:
         return [prompt] * max(1, prompt_count)
-    prompts = [
+    prompts: list[str] = [
         "\n\n".join(
             (
                 _YES_NO_PROMPT_ROOT,
diff --git a/tests/integration/test_sft_last_assistant_tokenization.py b/tests/integration/test_sft_last_assistant_tokenization.py
index 9b5e6e76b..49285c25c 100644
--- a/tests/integration/test_sft_last_assistant_tokenization.py
+++ b/tests/integration/test_sft_last_assistant_tokenization.py
@@ -2,6 +2,7 @@
 
 import pytest
 from transformers import AutoTokenizer
+from transformers.tokenization_utils_base import PreTrainedTokenizerBase
 
 from art.preprocessing.tokenize import tokenize_sft_batch
 from art.trajectories import Trajectory
@@ -36,6 +37,7 @@ def test_qwen_last_assistant_tool_call_tokenization() -> None:
         )
     except OSError:
         pytest.skip("Qwen3.6 tokenizer is not cached")
+    assert isinstance(tokenizer, PreTrainedTokenizerBase)
     tokenizer.chat_template = QWEN_LAST_ASSISTANT_TEMPLATE
     trajectory = Trajectory(
         messages_and_choices=[
@@ -94,6 +96,7 @@ def test_qwen_last_assistant_preserves_stock_generation_boundary() -> None:
         )
     except OSError:
         pytest.skip("Qwen3.6 tokenizer is not cached")
+    assert isinstance(tokenizer, PreTrainedTokenizerBase)
     trajectory = Trajectory(
         messages_and_choices=[
             {"role": "system", "content": "You are helpful."},
diff --git a/tests/unit/test_dedicated_config.py b/tests/unit/test_dedicated_config.py
index 5f09cfbce..be5d7c583 100644
--- a/tests/unit/test_dedicated_config.py
+++ b/tests/unit/test_dedicated_config.py
@@ -248,7 +248,7 @@ def test_invalid_rollout_weights_mode():
         ValueError, match="rollout_weights_mode must be either 'lora' or 'merged'"
     ):
         validate_dedicated_config(
-            InternalModelConfig(rollout_weights_mode="bad-mode")  # type: ignore[typeddict-item]
+            InternalModelConfig(rollout_weights_mode="bad-mode")  # type: ignore[typeddict-item]  # ty:ignore[invalid-argument-type]
         )
 
 
diff --git a/tests/unit/test_frontend_logging.py b/tests/unit/test_frontend_logging.py
index 8014c98b2..31a12ae2b 100644
--- a/tests/unit/test_frontend_logging.py
+++ b/tests/unit/test_frontend_logging.py
@@ -1264,7 +1264,7 @@ async def mock_train_model(*args, **kwargs):
                     TRAIN_GRADIENT_STEPS_KEY: 2.0,
                 }
 
-        backend._train_model = mock_train_model  # type: ignore[method-assign]
+        backend._train_model = mock_train_model  # type: ignore[method-assign]  # ty:ignore[invalid-assignment]
         backend._get_step = AsyncMock(return_value=1)  # type: ignore[method-assign]
 
         groups = [
diff --git a/tests/unit/test_moe_routing_real_path.py b/tests/unit/test_moe_routing_real_path.py
index 569824db5..244e6ef2b 100644
--- a/tests/unit/test_moe_routing_real_path.py
+++ b/tests/unit/test_moe_routing_real_path.py
@@ -141,7 +141,7 @@ def _tokenized(
         trajectory=Trajectory(),
         choice_offsets=[trainable_start],
         extra_logprobs={},
-        _tokenizer=_FakeTokenizer(),  # type: ignore[arg-type]
+        _tokenizer=_FakeTokenizer(),  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
         moe_routed_experts=cast(list[list[list[int]] | None], routes),
         prompt_id=prompt_id,
         prompt_length=prompt_length,
diff --git a/tests/unit/test_moe_routing_replay.py b/tests/unit/test_moe_routing_replay.py
index 3d944ace8..8d2338700 100644
--- a/tests/unit/test_moe_routing_replay.py
+++ b/tests/unit/test_moe_routing_replay.py
@@ -231,7 +231,7 @@ def __init__(self, router: _FakeRouter | None = None) -> None:
 
 def _fake_chunk_router(chunk: _FakeChunk) -> _FakeRouter:
     layer = cast(_FakeLayer, chunk.decoder.layers[0])
-    return cast(_FakeRouter, layer.mlp.router)
+    return layer.mlp.router
 
 
 def _assert_target(
diff --git a/tests/unit/test_pipeline_trainer_batching.py b/tests/unit/test_pipeline_trainer_batching.py
index 0ab412e8f..ba39e62a5 100644
--- a/tests/unit/test_pipeline_trainer_batching.py
+++ b/tests/unit/test_pipeline_trainer_batching.py
@@ -8,6 +8,10 @@
 from art.pipeline_trainer.trainer import PipelineTrainer
 
 
+async def _noop_rollout(*_args: object, **_kwargs: object) -> TrajectoryGroup:
+    return TrajectoryGroup([])
+
+
 def _make_group() -> TrajectoryGroup:
     return TrajectoryGroup(
         [
@@ -35,7 +39,7 @@ async def test_collect_batch_respects_max_batch_size(tmp_path: Path) -> None:
     trainer = PipelineTrainer(
         model=model,
         backend=MagicMock(),  # type: ignore[arg-type]
-        rollout_fn=lambda *_args, **_kwargs: asyncio.sleep(0),
+        rollout_fn=_noop_rollout,
         scenarios=[],
         config={},
         num_rollout_workers=1,
diff --git a/tests/unit/test_pipeline_trainer_local_backend.py b/tests/unit/test_pipeline_trainer_local_backend.py
index d8945fe62..32e57e0d5 100644
--- a/tests/unit/test_pipeline_trainer_local_backend.py
+++ b/tests/unit/test_pipeline_trainer_local_backend.py
@@ -27,6 +27,10 @@
 from art.utils.output_dirs import get_model_dir, get_step_checkpoint_dir
 
 
+async def _noop_rollout(*_args: object, **_kwargs: object) -> TrajectoryGroup:
+    return TrajectoryGroup([])
+
+
 def _make_group(rewards: list[float]) -> TrajectoryGroup:
     return TrajectoryGroup(
         [
@@ -51,8 +55,8 @@ def _make_trainer(
 ) -> PipelineTrainer:
     return PipelineTrainer(
         model=model,
-        backend=backend,  # type: ignore[arg-type]
-        rollout_fn=lambda *_args, **_kwargs: asyncio.sleep(0),
+        backend=backend,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+        rollout_fn=_noop_rollout,
         scenarios=[],
         config={},
         num_rollout_workers=1,
@@ -92,6 +96,7 @@ async def test_pipeline_trainer_preserves_backend_train_kwargs(tmp_path: Path) -
 
     await trainer._training_stage()
 
+    assert backend.train.await_args is not None
     assert backend.train.await_args.kwargs == {
         "learning_rate": 2e-5,
         "loss_fn": "cispo",
@@ -126,6 +131,7 @@ async def test_pipeline_trainer_forwards_default_kl_step_zero_for_generic_backen
 
     await trainer._training_stage()
 
+    assert backend.train.await_args is not None
     assert backend.train.await_args.kwargs == {
         "learning_rate": 1e-5,
         "loss_fn": "cispo",
@@ -166,6 +172,7 @@ async def test_pipeline_trainer_kl_step_lag_floors_at_zero(
 
     await trainer._training_stage()
 
+    assert backend.train.await_args is not None
     assert backend.train.await_args.kwargs["kl_penalty_reference_step"] == 0
 
 
@@ -196,6 +203,7 @@ async def test_pipeline_trainer_kl_step_lag_computes_reference(
 
     await trainer._training_stage()
 
+    assert backend.train.await_args is not None
     assert backend.train.await_args.kwargs["kl_penalty_reference_step"] == 1
 
 
@@ -245,7 +253,7 @@ async def test_pipeline_trainer_uses_same_train_kwargs_for_local_backend(
 
     await trainer._training_stage()
 
-    assert backend.train.await_args.kwargs == {  # type: ignore[attr-defined]
+    assert backend.train.await_args.kwargs == {  # type: ignore[attr-defined]  # ty:ignore[unresolved-attribute]
         "learning_rate": 3e-5,
         "loss_fn": "ppo",
         "loss_fn_config": None,
@@ -278,7 +286,7 @@ async def fake_train_model(
         seen["verbose"] = verbose
         yield {}
 
-    backend._train_model = fake_train_model  # type: ignore[method-assign]
+    backend._train_model = fake_train_model  # type: ignore[method-assign]  # ty:ignore[invalid-assignment]
     backend._get_step = AsyncMock(return_value=1)  # type: ignore[method-assign]
     with patch.object(model, "_get_wandb_run", return_value=None):
         result = await backend.train(
@@ -318,7 +326,7 @@ async def fake_train_model(
         seen["verbose"] = verbose
         yield {}
 
-    backend._train_model = fake_train_model  # type: ignore[method-assign]
+    backend._train_model = fake_train_model  # type: ignore[method-assign]  # ty:ignore[invalid-assignment]
     backend._get_step = AsyncMock(return_value=1)  # type: ignore[method-assign]
     with patch.object(model, "_get_wandb_run", return_value=None):
         result = await backend.train(
@@ -359,7 +367,7 @@ async def fake_train_model(
         seen["dev_config"] = dev_config
         yield {}
 
-    backend._train_model = fake_train_model  # type: ignore[method-assign]
+    backend._train_model = fake_train_model  # type: ignore[method-assign]  # ty:ignore[invalid-assignment]
     backend._get_step = AsyncMock(return_value=1)  # type: ignore[method-assign]
 
     with patch.object(model, "_get_wandb_run", return_value=None):
@@ -401,7 +409,7 @@ async def fake_train_model(
         seen["dev_config"] = dev_config
         yield {}
 
-    backend._train_model = fake_train_model  # type: ignore[method-assign]
+    backend._train_model = fake_train_model  # type: ignore[method-assign]  # ty:ignore[invalid-assignment]
     backend._get_step = AsyncMock(return_value=1)  # type: ignore[method-assign]
     with patch.object(model, "_get_wandb_run", return_value=None):
         await backend.train(
diff --git a/tests/unit/test_pipeline_trainer_metrics.py b/tests/unit/test_pipeline_trainer_metrics.py
index 86ad01cb3..bac2fd8da 100644
--- a/tests/unit/test_pipeline_trainer_metrics.py
+++ b/tests/unit/test_pipeline_trainer_metrics.py
@@ -10,6 +10,10 @@
 from art.pipeline_trainer.trainer import PipelineTrainer
 
 
+async def _noop_rollout(*_args: object, **_kwargs: object) -> TrajectoryGroup:
+    return TrajectoryGroup([])
+
+
 def _make_group(
     rewards: list[float], *, initial_policy_version: int | None
 ) -> TrajectoryGroup:
@@ -45,7 +49,7 @@ async def test_pipeline_trainer_logs_explicit_stale_and_zero_variance_metrics(
     trainer = PipelineTrainer(
         model=model,
         backend=backend,
-        rollout_fn=lambda *_args, **_kwargs: asyncio.sleep(0),
+        rollout_fn=_noop_rollout,
         scenarios=[],
         config={},
         num_rollout_workers=1,
diff --git a/tests/unit/test_preprocessing_tokenize.py b/tests/unit/test_preprocessing_tokenize.py
index 66aa2797b..46c6afc4c 100644
--- a/tests/unit/test_preprocessing_tokenize.py
+++ b/tests/unit/test_preprocessing_tokenize.py
@@ -181,7 +181,7 @@ def test_tokenize_sft_batch_masks_response_tokens_without_unsloth_import() -> No
     batch = tokenize_sft_batch(
         trajectory_batch=[Trajectory(messages_and_choices=messages, reward=1.0)],
         learning_rate=1e-5,
-        tokenizer=tokenizer,  # type: ignore[arg-type]
+        tokenizer=tokenizer,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
         instruction_part="<user>",
         response_part="<assistant>",
     )
@@ -207,7 +207,7 @@ def test_tokenize_sft_batch_all_trains_every_assistant_turn() -> None:
     batch = tokenize_sft_batch(
         trajectory_batch=[Trajectory(messages_and_choices=messages)],
         learning_rate=1e-5,
-        tokenizer=tokenizer,  # type: ignore[arg-type]
+        tokenizer=tokenizer,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
         instruction_part="<user>",
         response_part="<assistant>",
         assistant_turns="all",
@@ -231,7 +231,7 @@ def test_tokenize_sft_batch_passes_chat_template_kwargs() -> None:
     tokenize_sft_batch(
         trajectory_batch=[Trajectory(messages_and_choices=messages, reward=1.0)],
         learning_rate=1e-5,
-        tokenizer=tokenizer,  # type: ignore[arg-type]
+        tokenizer=tokenizer,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
         instruction_part="<user>",
         response_part="<assistant>",
         chat_template_kwargs={
@@ -269,7 +269,7 @@ def test_tokenize_sft_batch_trains_only_final_assistant_turn() -> None:
     batch = tokenize_sft_batch(
         trajectory_batch=[Trajectory(messages_and_choices=messages)],
         learning_rate=1e-5,
-        tokenizer=tokenizer,  # type: ignore[arg-type]
+        tokenizer=tokenizer,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
         instruction_part="<user>",
         response_part="<assistant>",
         assistant_turns="last",
@@ -313,7 +313,7 @@ def test_tokenize_sft_batch_trains_final_tool_call_and_turn_end() -> None:
     batch = tokenize_sft_batch(
         trajectory_batch=[Trajectory(messages_and_choices=messages)],
         learning_rate=1e-5,
-        tokenizer=tokenizer,  # type: ignore[arg-type]
+        tokenizer=tokenizer,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
         instruction_part="<user>",
         response_part="<assistant>",
         assistant_turns="last",
@@ -358,7 +358,7 @@ def test_tokenize_sft_batch_preserves_prompt_target_token_boundary() -> None:
     batch = tokenize_sft_batch(
         trajectory_batch=[Trajectory(messages_and_choices=messages)],
         learning_rate=1e-5,
-        tokenizer=tokenizer,  # type: ignore[arg-type]
+        tokenizer=tokenizer,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
         instruction_part="<user>",
         response_part="<assistant>",
         assistant_turns="last",
@@ -380,7 +380,7 @@ def test_tokenize_sft_batch_last_requires_final_assistant_message() -> None:
                 Trajectory(messages_and_choices=[{"role": "user", "content": "hello"}])
             ],
             learning_rate=1e-5,
-            tokenizer=tokenizer,  # type: ignore[arg-type]
+            tokenizer=tokenizer,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
             instruction_part="<user>",
             response_part="<assistant>",
             assistant_turns="last",
@@ -400,7 +400,7 @@ def test_tokenize_sft_batch_last_rejects_non_prefix_stable_template() -> None:
                 )
             ],
             learning_rate=1e-5,
-            tokenizer=tokenizer,  # type: ignore[arg-type]
+            tokenizer=tokenizer,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
             instruction_part="<user>",
             response_part="<assistant>",
             assistant_turns="last",
diff --git a/tests/unit/test_serverless_pipeline_trainer_compat.py b/tests/unit/test_serverless_pipeline_trainer_compat.py
index 4f3e8a3e7..7a083fd21 100644
--- a/tests/unit/test_serverless_pipeline_trainer_compat.py
+++ b/tests/unit/test_serverless_pipeline_trainer_compat.py
@@ -56,7 +56,7 @@ async def fake_train_model(
         seen["verbose"] = verbose
         yield {"loss": 0.25}
 
-    backend._train_model = fake_train_model  # type: ignore[method-assign]
+    backend._train_model = fake_train_model  # type: ignore[method-assign]  # ty:ignore[invalid-assignment]
     backend._get_step = AsyncMock(return_value=3)  # type: ignore[method-assign]
 
     with patch.object(model, "_get_wandb_run", return_value=None):
@@ -156,7 +156,7 @@ async def test_serverless_train_model_forwards_experimental_config() -> None:
     async def events_list(**_kwargs: Any):
         yield SimpleNamespace(id="event-id", type="training_ended", data={})
 
-    backend._client.training_jobs.events.list = events_list  # type: ignore[attr-defined]
+    backend._client.training_jobs.events.list = events_list  # type: ignore[attr-defined]  # ty:ignore[invalid-assignment]
 
     async def no_sleep(_seconds: float) -> None:
         return None
@@ -227,7 +227,7 @@ async def test_serverless_train_sft_forwards_metric_logging_config() -> None:
     async def events_list(**_kwargs: Any):
         yield SimpleNamespace(id="event-id", type="training_ended", data={})
 
-    backend._client.sft_training_jobs.events.list = events_list  # type: ignore[attr-defined]
+    backend._client.sft_training_jobs.events.list = events_list  # type: ignore[attr-defined]  # ty:ignore[invalid-assignment]
 
     async def no_sleep(_seconds: float) -> None:
         return None
@@ -332,7 +332,7 @@ async def fake_train_model(
         seen["dev_config"] = dev_config
         yield {}
 
-    backend._train_model = fake_train_model  # type: ignore[method-assign]
+    backend._train_model = fake_train_model  # type: ignore[method-assign]  # ty:ignore[invalid-assignment]
     backend._get_step = AsyncMock(return_value=1)  # type: ignore[method-assign]
 
     with patch.object(model, "_get_wandb_run", return_value=None):
diff --git a/tests/unit/test_tinker_renderers.py b/tests/unit/test_tinker_renderers.py
index 59b01961b..87f3d4422 100644
--- a/tests/unit/test_tinker_renderers.py
+++ b/tests/unit/test_tinker_renderers.py
@@ -53,7 +53,7 @@ def decode(self, tokens: int | list[int]) -> str:
 
 def _decode_model_input(tokenizer: FakeTokenizer, model_input: object) -> str:
     tokens: list[int] = []
-    for chunk in model_input.chunks:  # type: ignore[attr-defined]
+    for chunk in model_input.chunks:  # type: ignore[attr-defined]  # ty:ignore[unresolved-attribute]
         assert hasattr(chunk, "tokens"), f"Unexpected non-text chunk: {chunk!r}"
         tokens.extend(list(chunk.tokens))
     return tokenizer.decode(tokens)
diff --git a/tests/unit/test_tokenize_trajectory_groups.ipynb b/tests/unit/test_tokenize_trajectory_groups.ipynb
index 7f19af951..9e80d1a48 100644
--- a/tests/unit/test_tokenize_trajectory_groups.ipynb
+++ b/tests/unit/test_tokenize_trajectory_groups.ipynb
@@ -37,11 +37,13 @@
     "from openai.types.chat.chat_completion_message import ChatCompletionMessage\n",
     "from openai.types.chat.chat_completion_token_logprob import ChatCompletionTokenLogprob\n",
     "from transformers import AutoTokenizer\n",
+    "from transformers.tokenization_utils_base import PreTrainedTokenizerBase\n",
     "\n",
     "import art\n",
     "from art.preprocessing.tokenize import tokenize_trajectory_groups\n",
     "\n",
-    "tokenizer = AutoTokenizer.from_pretrained(\"Qwen/Qwen2.5-7B-Instruct\")"
+    "tokenizer = AutoTokenizer.from_pretrained(\"Qwen/Qwen2.5-7B-Instruct\")\n",
+    "assert isinstance(tokenizer, PreTrainedTokenizerBase)"
    ]
   },
   {
diff --git a/tests/unit/test_track_api_cost.py b/tests/unit/test_track_api_cost.py
index ed1d8f37b..dc4969241 100644
--- a/tests/unit/test_track_api_cost.py
+++ b/tests/unit/test_track_api_cost.py
@@ -715,7 +715,7 @@ async def eval_fn(
         trainer = PipelineTrainer(
             model=model,
             backend=backend,
-            rollout_fn=lambda *_args, **_kwargs: asyncio.sleep(0),
+            rollout_fn=lambda *_args, **_kwargs: asyncio.sleep(0),  # ty: ignore[invalid-argument-type]
             scenarios=[],
             config={},
             num_rollout_workers=1,
diff --git a/tests/unit/test_trainer_rank_validation.py b/tests/unit/test_trainer_rank_validation.py
index 5bf518f11..8dfbe88b5 100644
--- a/tests/unit/test_trainer_rank_validation.py
+++ b/tests/unit/test_trainer_rank_validation.py
@@ -110,7 +110,7 @@ def _output_values(outputs: object) -> list[int]:
         assert isinstance(target_logprobs, torch.Tensor)
         return [int(target_logprobs.item())]
     values: list[int] = []
-    for item in outputs:  # type: ignore[union-attr]
+    for item in outputs:  # type: ignore[union-attr]  # ty:ignore[not-iterable]
         values.extend(_output_values(item))
     return values
 
@@ -118,14 +118,14 @@ def _output_values(outputs: object) -> list[int]:
 def _output_shape(outputs: object) -> object:
     if isinstance(outputs, ForwardOutput):
         return "output"
-    return [_output_shape(item) for item in outputs]  # type: ignore[union-attr]
+    return [_output_shape(item) for item in outputs]  # type: ignore[union-attr]  # ty:ignore[not-iterable]
 
 
 def _trainer_with_checkpoint(
     monkeypatch: pytest.MonkeyPatch,
     value: torch.Tensor,
 ) -> tuple[TrainerRank, torch.nn.Parameter]:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     param = torch.nn.Parameter(value.clone())
     trainer._checkpoint_slot_params_by_name["student"] = (param,)
     monkeypatch.setattr(
@@ -140,7 +140,7 @@ def _tracked_targets(
     trainer: TrainerRank, ref: _SlotRef, *scales: float
 ) -> list[torch.Tensor]:
     tracked = trainer._track_slot_graph_outputs(
-        ref,  # type: ignore[arg-type]
+        ref,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
         [
             ForwardOutput(torch.ones(1, requires_grad=True) * scale, None, None, None)
             for scale in scales
@@ -155,14 +155,14 @@ def test_forward_input_validation() -> None:
     with pytest.raises(ValueError, match="cannot set both checkpoint and lora"):
         ForwardInput(input_tokens=torch.tensor([1]), checkpoint="a", lora="b")
     with pytest.raises(ValueError, match="top_k=9 exceeds vocabulary size 8"):
-        _validate_top_k(9, _Model())  # type: ignore[arg-type]
+        _validate_top_k(9, _Model())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
 
 
 @pytest.mark.parametrize(("checkpoint", "expected"), ((Unset, Unset), (None, None)))
 def test_forward_input_distinguishes_unset_and_base_checkpoint(
     checkpoint: object, expected: object
 ) -> None:
-    request = ForwardInput(input_tokens=torch.tensor([1]), checkpoint=checkpoint)  # type: ignore[arg-type]
+    request = ForwardInput(input_tokens=torch.tensor([1]), checkpoint=checkpoint)  # type: ignore[arg-type]  # ty:ignore[no-matching-overload]
 
     assert request.checkpoint is expected
     assert request.lora is Unset
@@ -176,24 +176,24 @@ def test_forward_input_preserves_public_runtime_shape() -> None:
 
 @pytest.mark.parametrize("depth", (0, 2))
 def test_trainer_rank_accepts_shared_prefix_depth(depth: int) -> None:
-    trainer = TrainerRank(_runtime(), shared_prefix_max_depth=depth)  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime(), shared_prefix_max_depth=depth)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
 
     assert trainer.shared_prefix_max_depth == depth
 
 
 def test_trainer_rank_adapter_stack_errors() -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
 
     with pytest.raises(RuntimeError, match="No pushed LoRA or checkpoint"):
         trainer.pop_pushed_lora_or_checkpoint()
-    trainer._slot_stack.append(object())  # type: ignore[arg-type]
+    trainer._slot_stack.append(object())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     for load in (trainer.load_checkpoint_slot, trainer.load_lora_slot):
         with pytest.raises(RuntimeError, match="Cannot load a LoRA/checkpoint"):
             load("teacher", {})
 
 
 def test_trainer_rank_rejects_adapter_keys_without_installed_lora_site() -> None:
-    trainer = TrainerRank(_runtime(_FakeLoRASite("base.layer")))  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime(_FakeLoRASite("base.layer")))  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     valid = {
         "base.layer.lora_A.weight": torch.empty(1),
         "base.layer.lora_B.weight": torch.empty(1),
@@ -210,7 +210,7 @@ def test_trainer_rank_rejects_adapter_keys_without_installed_lora_site() -> None
 
 def test_trainer_rank_normalizes_adapter_tensors_to_installed_site() -> None:
     site = _FakeLoRASite("base.layer", dtype=torch.bfloat16)
-    trainer = TrainerRank(_runtime(site))  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime(site))  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     adapter = {
         "base.layer.lora_A.weight": torch.ones(3, 4, dtype=torch.float32),
         "base.layer.lora_B.weight": torch.ones(5, 3, dtype=torch.float32),
@@ -223,7 +223,7 @@ def test_trainer_rank_normalizes_adapter_tensors_to_installed_site() -> None:
 
 
 def test_checkpoint_slot_adapter_config_is_validated_and_copied() -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     config = {
         "base_model_name_or_path": "Qwen/Qwen3-8B",
         "r": 8,
@@ -250,7 +250,7 @@ def test_checkpoint_slot_adapter_config_is_validated_and_copied() -> None:
 def test_checkpoint_slot_adapter_config_rejects_cross_rank_mismatch(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     monkeypatch.setattr("art.trainer_rank.dist.is_initialized", lambda: True)
     monkeypatch.setattr("art.trainer_rank.dist.get_world_size", lambda: 2)
 
@@ -266,7 +266,7 @@ def gather(output: list[object], value: object) -> None:
 def test_load_checkpoint_slot_retains_config_and_uses_its_alpha(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     seen: dict[str, object] = {}
     monkeypatch.setattr(
         trainer,
@@ -294,7 +294,7 @@ def test_load_checkpoint_slot_retains_config_and_uses_its_alpha(
 
 
 def test_checkpoint_slot_publish_requires_retained_adapter_config() -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     with pytest.raises(ValueError, match="Unknown checkpoint slot"):
         trainer.save_checkpoint_slot_lora("missing", "/unused")
 
@@ -305,7 +305,7 @@ def test_checkpoint_slot_publish_requires_retained_adapter_config() -> None:
 
 
 def test_trainer_rank_default_forward_uses_explicit_base_slot() -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
 
     plan = trainer._plan_flat_forward([_target_request(1)])
 
@@ -318,7 +318,7 @@ def test_trainer_rank_default_forward_uses_explicit_base_slot() -> None:
 
 def test_optim_step_requires_loaded_checkpoint_slot() -> None:
     optimizer = _NativeOptimizer()
-    trainer = TrainerRank(_runtime(optimizer=optimizer))  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime(optimizer=optimizer))  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
 
     with pytest.raises(TrainerRankSlotStateError, match="loaded checkpoint slot"):
         trainer.optim_step(params=AdamParams(learning_rate=1e-3))
@@ -327,7 +327,7 @@ def test_optim_step_requires_loaded_checkpoint_slot() -> None:
 
 
 def test_optim_step_rejects_loaded_slots_without_grads() -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     trainer._checkpoint_slot_params_by_name["student"] = (
         torch.nn.Parameter(torch.ones(2)),
     )
@@ -344,7 +344,7 @@ def test_optim_step_rejects_loaded_slots_without_grads() -> None:
 def test_optim_step_rejects_explicit_slot_subset_with_missing_grads(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     ready = torch.nn.Parameter(torch.ones(2))
     missing = torch.nn.Parameter(torch.ones(2))
     ready.grad = torch.ones_like(ready)
@@ -366,7 +366,7 @@ def test_optim_step_rejects_explicit_slot_subset_with_missing_grads(
 def test_optim_step_implicitly_steps_only_slots_with_grads(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     ready = torch.nn.Parameter(torch.ones(2))
     untouched = torch.nn.Parameter(torch.ones(2))
     ready.grad = torch.ones_like(ready)
@@ -488,12 +488,12 @@ def test_trainer_rank_rejects_mutating_slot_with_pending_graph(
     operation: str,
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     ref = _SlotRef("checkpoint", "teacher")
     monkeypatch.setattr(trainer, "_slot_ref", lambda kind, name: _SlotRef(kind, name))
     target = _tracked_targets(trainer, ref, 2)[0]
     guard = (
-        (lambda: trainer._guard_slot_can_load(ref))  # type: ignore[arg-type]
+        (lambda: trainer._guard_slot_can_load(ref))  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
         if operation == "load"
         else (lambda: trainer._guard_checkpoint_can_step("teacher"))
     )
@@ -515,7 +515,7 @@ def test_trainer_rank_step_allows_missing_slot_graph_bookkeeping(
 
 
 def test_trainer_rank_zero_grad_does_not_clear_live_slot_graphs() -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     ref = _SlotRef("lora", "teacher")
     output = ForwardOutput(
         None,
@@ -527,42 +527,42 @@ def test_trainer_rank_zero_grad_does_not_clear_live_slot_graphs() -> None:
         None,
     )
 
-    tracked = trainer._track_slot_graph_outputs(ref, [output])  # type: ignore[arg-type]
+    tracked = trainer._track_slot_graph_outputs(ref, [output])  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     trainer.zero_grad()
 
     assert tracked[0].top_k is not None
     with pytest.raises(TrainerRankSlotStateError, match="live backward graph"):
-        trainer._guard_slot_can_load(ref)  # type: ignore[arg-type]
+        trainer._guard_slot_can_load(ref)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
 
 
 def test_trainer_rank_retained_backward_keeps_slot_graph_guard() -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     ref = _SlotRef("checkpoint", "teacher")
     target = _tracked_targets(trainer, ref, 2)[0]
 
     target.sum().backward(retain_graph=True)
     with pytest.raises(TrainerRankSlotStateError, match="live backward graph"):
-        trainer._guard_slot_can_load(ref)  # type: ignore[arg-type]
+        trainer._guard_slot_can_load(ref)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
 
     target.sum().backward()
-    trainer._guard_slot_can_load(ref)  # type: ignore[arg-type]
+    trainer._guard_slot_can_load(ref)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
 
 
 def test_trainer_rank_tracks_each_independent_output_graph() -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     ref = _SlotRef("checkpoint", "teacher")
     first, second = _tracked_targets(trainer, ref, 2, 3)
 
     first.sum().backward()
     with pytest.raises(TrainerRankSlotStateError, match="live backward graph"):
-        trainer._guard_slot_can_load(ref)  # type: ignore[arg-type]
+        trainer._guard_slot_can_load(ref)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
 
     second.sum().backward()
-    trainer._guard_slot_can_load(ref)  # type: ignore[arg-type]
+    trainer._guard_slot_can_load(ref)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
 
 
 def test_trainer_rank_tracks_graph_after_output_is_replaced_by_loss() -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     ref = _SlotRef("checkpoint", "teacher")
     target = _tracked_targets(trainer, ref, 2)[0]
     loss = target.sum()
@@ -570,24 +570,24 @@ def test_trainer_rank_tracks_graph_after_output_is_replaced_by_loss() -> None:
     gc.collect()
 
     with pytest.raises(TrainerRankSlotStateError, match="live backward graph"):
-        trainer._guard_slot_can_load(ref)  # type: ignore[arg-type]
+        trainer._guard_slot_can_load(ref)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
 
     loss.backward()
-    trainer._guard_slot_can_load(ref)  # type: ignore[arg-type]
+    trainer._guard_slot_can_load(ref)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
 
 
 def test_trainer_rank_releases_abandoned_output_graph() -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     ref = _SlotRef("checkpoint", "teacher")
     target = _tracked_targets(trainer, ref, 2)[0]
     del target
     gc.collect()
 
-    trainer._guard_slot_can_load(ref)  # type: ignore[arg-type]
+    trainer._guard_slot_can_load(ref)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
 
 
 def test_dp_rank_forward_preserves_nested_shape_for_inactive_requests() -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     request_a = ForwardInput(input_tokens=torch.tensor([1]))
     request_b = ForwardInput(input_tokens=torch.tensor([2]))
 
@@ -604,7 +604,7 @@ def test_dp_rank_forward_preserves_nested_shape_for_inactive_requests() -> None:
 def test_dp_rank_forward_supports_arbitrary_nested_depth(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     _stub_forward(monkeypatch, trainer, _indexed_outputs)
     nested = [
         [[[[[_target_request(1)]]]]],
@@ -623,7 +623,7 @@ def test_dp_rank_forward_supports_arbitrary_nested_depth(
 def test_forward_micro_batches_uses_deterministic_dp_windows(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     _stub_forward(monkeypatch, trainer, dp=(1, 2))
 
     batches = list(
@@ -637,7 +637,7 @@ def test_forward_micro_batches_uses_deterministic_dp_windows(
 def test_forward_micro_batches_syncs_fit_decision_across_dp(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     _stub_forward(monkeypatch, trainer, dp=(1, 2), profiled=True)
     sync_flags: list[bool] = []
 
@@ -659,7 +659,7 @@ def memory_check(required: int, *, sync_across_dp: bool = False) -> _MemoryCheck
 def test_forward_micro_batches_supports_arbitrary_nested_depth(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     _stub_forward(monkeypatch, trainer, _indexed_outputs, profiled=True)
     expected = [
         [[[[[_target_request(1)]]]]],
@@ -680,7 +680,7 @@ def test_forward_micro_batches_supports_arbitrary_nested_depth(
 def test_forward_micro_batches_ramps_after_first_success(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
 
     def run(plan, **_kwargs):
         trainer._memory_profiles[plan.signature] = _MemoryProfile(
@@ -706,7 +706,7 @@ def run(plan, **_kwargs):
 def test_forward_micro_batches_does_not_overtrust_tiny_memory_profile(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     monkeypatch.setattr(trainer, "_dp_rank_and_size", lambda: (0, 1))
     inputs = [_target_request(i) for i in range(64)]
     tiny_plan = trainer._plan_flat_forward([inputs[0]])
@@ -724,7 +724,7 @@ def test_forward_micro_batches_does_not_overtrust_tiny_memory_profile(
 def test_forward_micro_batches_tail_does_not_reset_stable_window(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     trainer._last_global_micro_batch_size = 64
     _stub_forward(monkeypatch, trainer, profiled=True)
     monkeypatch.setattr(
@@ -752,7 +752,7 @@ def test_forward_micro_batches_tail_does_not_reset_stable_window(
 def test_forward_micro_batches_raises_when_smallest_batch_will_not_fit(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     monkeypatch.setattr(trainer, "_dp_rank_and_size", lambda: (0, 1))
     monkeypatch.setattr(
         trainer,
@@ -775,7 +775,7 @@ def test_forward_micro_batches_raises_when_smallest_batch_will_not_fit(
 def test_forward_micro_batches_rejects_mismatched_replicated_counts(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     import art.trainer_rank as trainer_rank
 
     monkeypatch.setattr(trainer_rank.dist, "is_available", lambda: True)
@@ -814,7 +814,7 @@ def __init__(self) -> None:
         def _preprocess(self, *args: object, **kwargs: object) -> None:
             return None
 
-    trainer = TrainerRank(_runtime(FakeGPT()))  # type: ignore[arg-type]
+    trainer = TrainerRank(_runtime(FakeGPT()))  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     tokens = torch.tensor([[1, 2, 3]], dtype=torch.long)
     labels = torch.stack((tokens, tokens + 1), dim=-1)
 
diff --git a/tests/unit/test_trainer_rank_weird_shapes.py b/tests/unit/test_trainer_rank_weird_shapes.py
index 841169049..a5562db01 100644
--- a/tests/unit/test_trainer_rank_weird_shapes.py
+++ b/tests/unit/test_trainer_rank_weird_shapes.py
@@ -89,11 +89,11 @@ def _set_packed_token_budget(
     monkeypatch.setattr(
         rank,
         "_estimate_required_memory_bytes_from_values",
-        lambda **kwargs: kwargs["packed_tokens"],
+        lambda *, packed_tokens, **_kwargs: packed_tokens,
     )
 
     def check(required: int, *, sync_across_dp: bool = False) -> _MemoryCheck:
-        limit = available() if callable(available) else available
+        limit = available if isinstance(available, int) else available()
         return _MemoryCheck(required, limit, required <= limit)
 
     monkeypatch.setattr(rank, "_memory_check_required", check)
@@ -187,10 +187,14 @@ def test_shared_trainable_tokens_accumulate_independent_output_gradients() -> No
     hidden = torch.randn(int(pack.tokens.numel()), 3, requires_grad=True)
     weights = (2.0, 5.0)
 
-    loss = sum(
-        weight * hidden.index_select(0, positions).sum()
-        for weight, positions in zip(weights, pack.positions_by_sequence, strict=True)
-    )
+    loss = torch.stack(
+        [
+            weight * hidden.index_select(0, positions).sum()
+            for weight, positions in zip(
+                weights, pack.positions_by_sequence, strict=True
+            )
+        ]
+    ).sum()
     loss.backward()
 
     expected = torch.zeros_like(hidden)
@@ -204,7 +208,7 @@ def test_shared_trainable_tokens_accumulate_independent_output_gradients() -> No
 
 
 def test_planner_handles_vineppo_nested_shape_and_request_mix() -> None:
-    rank = TrainerRank(_runtime(), shared_prefix_max_depth=3)  # type: ignore[arg-type]
+    rank = TrainerRank(_runtime(), shared_prefix_max_depth=3)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     inputs = _vineppo_like_inputs()
     flat = list(_flatten(inputs))
 
@@ -227,7 +231,7 @@ def test_planner_handles_vineppo_nested_shape_and_request_mix() -> None:
 def test_forward_micro_batches_preserves_nested_vineppo_groups(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    rank = TrainerRank(_runtime(), shared_prefix_max_depth=2)  # type: ignore[arg-type]
+    rank = TrainerRank(_runtime(), shared_prefix_max_depth=2)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1))
     monkeypatch.setattr(rank, "_all_ranks_have_memory_profile", lambda **_kwargs: True)
     monkeypatch.setattr(
@@ -260,7 +264,7 @@ def test_forward_micro_batches_preserves_nested_vineppo_groups(
 def test_adaptive_planner_materializes_only_final_large_candidate(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    rank = TrainerRank(_runtime(), shared_prefix_max_depth=3)  # type: ignore[arg-type]
+    rank = TrainerRank(_runtime(), shared_prefix_max_depth=3)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     rank._last_global_micro_batch_size = 32
     monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1))
     monkeypatch.setattr(rank, "_all_ranks_have_memory_profile", lambda **_kwargs: True)
@@ -306,7 +310,7 @@ def estimate(requests):
 def test_adaptive_planner_globally_falls_back_when_one_rank_cannot_estimate(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    rank = TrainerRank(_runtime())  # type: ignore[arg-type]
+    rank = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 2))
     monkeypatch.setattr(rank, "_all_ranks_true", lambda _local: False)
     plans = 0
@@ -329,7 +333,7 @@ def plan(requests):
 def test_adaptive_planner_probes_new_heterogeneous_signatures(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    rank = TrainerRank(_runtime())  # type: ignore[arg-type]
+    rank = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1))
     monkeypatch.setattr(rank, "_resolve_slot_ref", lambda request: request.checkpoint)
     inputs = [
@@ -359,7 +363,7 @@ def test_adaptive_planner_probes_new_heterogeneous_signatures(
 def test_adaptive_planner_grows_stable_window_to_largest_aligned_fit(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    rank = TrainerRank(_runtime(), shared_prefix_max_depth=1)  # type: ignore[arg-type]
+    rank = TrainerRank(_runtime(), shared_prefix_max_depth=1)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     rank._last_global_micro_batch_size = 512
     monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1))
     monkeypatch.setattr(rank, "_all_ranks_have_memory_profile", lambda **_kwargs: True)
@@ -377,7 +381,7 @@ def test_adaptive_planner_grows_stable_window_to_largest_aligned_fit(
 def test_forward_micro_batches_shrinks_when_memory_budget_drops(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    rank = TrainerRank(_runtime(), shared_prefix_max_depth=2)  # type: ignore[arg-type]
+    rank = TrainerRank(_runtime(), shared_prefix_max_depth=2)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1))
     monkeypatch.setattr(rank, "_all_ranks_have_memory_profile", lambda **_kwargs: True)
     inputs = [_target_request(_tokens(1, 2, 3, index)) for index in range(14)]
@@ -426,7 +430,7 @@ def run(plan, **_kwargs):
 def test_heterogeneous_slots_split_packing_without_losing_output_estimates(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    rank = TrainerRank(_runtime(), shared_prefix_max_depth=4)  # type: ignore[arg-type]
+    rank = TrainerRank(_runtime(), shared_prefix_max_depth=4)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     monkeypatch.setattr(
         TrainerRank,
         "_slot_ref",
@@ -462,7 +466,7 @@ def test_forward_raises_before_expected_oom_with_actionable_context(
     api: str,
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    rank = TrainerRank(_runtime())  # type: ignore[arg-type]
+    rank = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     if api == "dp_rank_forward":
         monkeypatch.setattr(
             rank,
@@ -504,7 +508,7 @@ def test_flatten_rejects_dicts_to_avoid_silent_top_level_shape_changes() -> None
 
 
 def test_no_output_requests_do_not_pack_or_consume_compute_memory() -> None:
-    rank = TrainerRank(_runtime())  # type: ignore[arg-type]
+    rank = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
     requests: Iterable[ForwardInput] = [
         ForwardInput(input_tokens=_tokens(1, 2, 3)),
         ForwardInput(input_tokens=_tokens(1, 2, 4)),
diff --git a/uv.lock b/uv.lock
index 2486ef25d..bd7c959f7 100644
--- a/uv.lock
+++ b/uv.lock
@@ -4933,7 +4933,7 @@ dev = [
     { name = "pytest-xdist", specifier = ">=3.8.0" },
     { name = "ruff", specifier = ">=0.12.1" },
     { name = "skypilot", extras = ["cudo", "do", "fluidstack", "gcp", "kubernetes", "lambda", "paperspace", "runpod"], specifier = "==0.11.1" },
-    { name = "ty", specifier = "==0.0.14" },
+    { name = "ty", specifier = "==0.0.59" },
     { name = "uv", specifier = ">=0.11.7" },
 ]
 
@@ -8269,26 +8269,27 @@ wheels = [
 
 [[package]]
 name = "ty"
-version = "0.0.14"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/af/57/22c3d6bf95c2229120c49ffc2f0da8d9e8823755a1c3194da56e51f1cc31/ty-0.0.14.tar.gz", hash = "sha256:a691010565f59dd7f15cf324cdcd1d9065e010c77a04f887e1ea070ba34a7de2", size = 5036573, upload-time = "2026-01-27T00:57:31.427Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/99/cb/cc6d1d8de59beb17a41f9a614585f884ec2d95450306c173b3b7cc090d2e/ty-0.0.14-py3-none-linux_armv6l.whl", hash = "sha256:32cf2a7596e693094621d3ae568d7ee16707dce28c34d1762947874060fdddaa", size = 10034228, upload-time = "2026-01-27T00:57:53.133Z" },
-    { url = "https://files.pythonhosted.org/packages/f3/96/dd42816a2075a8f31542296ae687483a8d047f86a6538dfba573223eaf9a/ty-0.0.14-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:f971bf9805f49ce8c0968ad53e29624d80b970b9eb597b7cbaba25d8a18ce9a2", size = 9939162, upload-time = "2026-01-27T00:57:43.857Z" },
-    { url = "https://files.pythonhosted.org/packages/ff/b4/73c4859004e0f0a9eead9ecb67021438b2e8e5fdd8d03e7f5aca77623992/ty-0.0.14-py3-none-macosx_11_0_arm64.whl", hash = "sha256:45448b9e4806423523268bc15e9208c4f3f2ead7c344f615549d2e2354d6e924", size = 9418661, upload-time = "2026-01-27T00:58:03.411Z" },
-    { url = "https://files.pythonhosted.org/packages/58/35/839c4551b94613db4afa20ee555dd4f33bfa7352d5da74c5fa416ffa0fd2/ty-0.0.14-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ee94a9b747ff40114085206bdb3205a631ef19a4d3fb89e302a88754cbbae54c", size = 9837872, upload-time = "2026-01-27T00:57:23.718Z" },
-    { url = "https://files.pythonhosted.org/packages/41/2b/bbecf7e2faa20c04bebd35fc478668953ca50ee5847ce23e08acf20ea119/ty-0.0.14-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:6756715a3c33182e9ab8ffca2bb314d3c99b9c410b171736e145773ee0ae41c3", size = 9848819, upload-time = "2026-01-27T00:57:58.501Z" },
-    { url = "https://files.pythonhosted.org/packages/be/60/3c0ba0f19c0f647ad9d2b5b5ac68c0f0b4dc899001bd53b3a7537fb247a2/ty-0.0.14-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:89d0038a2f698ba8b6fec5cf216a4e44e2f95e4a5095a8c0f57fe549f87087c2", size = 10324371, upload-time = "2026-01-27T00:57:29.291Z" },
-    { url = "https://files.pythonhosted.org/packages/24/32/99d0a0b37d0397b0a989ffc2682493286aa3bc252b24004a6714368c2c3d/ty-0.0.14-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:2c64a83a2d669b77f50a4957039ca1450626fb474619f18f6f8a3eb885bf7544", size = 10865898, upload-time = "2026-01-27T00:57:33.542Z" },
-    { url = "https://files.pythonhosted.org/packages/1a/88/30b583a9e0311bb474269cfa91db53350557ebec09002bfc3fb3fc364e8c/ty-0.0.14-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:242488bfb547ef080199f6fd81369ab9cb638a778bb161511d091ffd49c12129", size = 10555777, upload-time = "2026-01-27T00:58:05.853Z" },
-    { url = "https://files.pythonhosted.org/packages/cd/a2/cb53fb6325dcf3d40f2b1d0457a25d55bfbae633c8e337bde8ec01a190eb/ty-0.0.14-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4790c3866f6c83a4f424fc7d09ebdb225c1f1131647ba8bdc6fcdc28f09ed0ff", size = 10412913, upload-time = "2026-01-27T00:57:38.834Z" },
-    { url = "https://files.pythonhosted.org/packages/42/8f/f2f5202d725ed1e6a4e5ffaa32b190a1fe70c0b1a2503d38515da4130b4c/ty-0.0.14-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:950f320437f96d4ea9a2332bbfb5b68f1c1acd269ebfa4c09b6970cc1565bd9d", size = 9837608, upload-time = "2026-01-27T00:57:55.898Z" },
-    { url = "https://files.pythonhosted.org/packages/f7/ba/59a2a0521640c489dafa2c546ae1f8465f92956fede18660653cce73b4c5/ty-0.0.14-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:4a0ec3ee70d83887f86925bbc1c56f4628bd58a0f47f6f32ddfe04e1f05466df", size = 9884324, upload-time = "2026-01-27T00:57:46.786Z" },
-    { url = "https://files.pythonhosted.org/packages/03/95/8d2a49880f47b638743212f011088552ecc454dd7a665ddcbdabea25772a/ty-0.0.14-py3-none-musllinux_1_2_i686.whl", hash = "sha256:a1a4e6b6da0c58b34415955279eff754d6206b35af56a18bb70eb519d8d139ef", size = 10033537, upload-time = "2026-01-27T00:58:01.149Z" },
-    { url = "https://files.pythonhosted.org/packages/e9/40/4523b36f2ce69f92ccf783855a9e0ebbbd0f0bb5cdce6211ee1737159ed3/ty-0.0.14-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:dc04384e874c5de4c5d743369c277c8aa73d1edea3c7fc646b2064b637db4db3", size = 10495910, upload-time = "2026-01-27T00:57:26.691Z" },
-    { url = "https://files.pythonhosted.org/packages/08/d5/655beb51224d1bfd4f9ddc0bb209659bfe71ff141bcf05c418ab670698f0/ty-0.0.14-py3-none-win32.whl", hash = "sha256:b20e22cf54c66b3e37e87377635da412d9a552c9bf4ad9fc449fed8b2e19dad2", size = 9507626, upload-time = "2026-01-27T00:57:41.43Z" },
-    { url = "https://files.pythonhosted.org/packages/b6/d9/c569c9961760e20e0a4bc008eeb1415754564304fd53997a371b7cf3f864/ty-0.0.14-py3-none-win_amd64.whl", hash = "sha256:e312ff9475522d1a33186657fe74d1ec98e4a13e016d66f5758a452c90ff6409", size = 10437980, upload-time = "2026-01-27T00:57:36.422Z" },
-    { url = "https://files.pythonhosted.org/packages/ad/0c/186829654f5bfd9a028f6648e9caeb11271960a61de97484627d24443f91/ty-0.0.14-py3-none-win_arm64.whl", hash = "sha256:b6facdbe9b740cb2c15293a1d178e22ffc600653646452632541d01c36d5e378", size = 9885831, upload-time = "2026-01-27T00:57:49.747Z" },
+version = "0.0.59"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/95/b0/84ae7b3bf6e3e9f57eb9635eeff5a80b36e57aa089f40be0fb5c384fa176/ty-0.0.59.tar.gz", hash = "sha256:53e53ffeed78ad59cd237fa8ea1316d2b94e13efdea9a945698acab549e005aa", size = 6145435, upload-time = "2026-07-12T20:22:02.781Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/57/e8/650b42fbef4d48e6ca682b0b6e9b68fa8fcf55cbb0a6892ab89990018b6f/ty-0.0.59-py3-none-linux_armv6l.whl", hash = "sha256:f8fb08a767ef8f11ea3c537b9d77860726cc2bc39e6f77ad13c02d5b289f20a7", size = 11700328, upload-time = "2026-07-12T20:21:26.046Z" },
+    { url = "https://files.pythonhosted.org/packages/22/ac/0ca3a89d5f59ae5f308e5e83428cac5f9143200767743e052fba90b4b81e/ty-0.0.59-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:c7f4d5630836c8a0ba13dd4ac7bdae080a7d6ebe965b817ff642dc961bcf2a53", size = 11494310, upload-time = "2026-07-12T20:21:28.491Z" },
+    { url = "https://files.pythonhosted.org/packages/f3/f8/5076de6001cefbccd8e6dc8472262697e43308ff66b0e87c72abba136357/ty-0.0.59-py3-none-macosx_11_0_arm64.whl", hash = "sha256:872f6fb02c6db5553c4d5fb283b3d50f0985fb9a29a910e4fda4793a775c1926", size = 11026797, upload-time = "2026-07-12T20:21:30.879Z" },
+    { url = "https://files.pythonhosted.org/packages/2e/0f/fca28481b6a138e2b798ad9fdc98a095475f9104948ba242fce4b477782b/ty-0.0.59-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2af8eefbfe806337770eec12c0c819c5f1b8f5b85f8369cb1cc9fa25234a2208", size = 11475304, upload-time = "2026-07-12T20:21:33.041Z" },
+    { url = "https://files.pythonhosted.org/packages/08/4b/1fed8b81b389ef4bbc0400f19e05fc16496b162577779dc0e5fc65ac216c/ty-0.0.59-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:0acf8b76a1c9a7ddef460b42475f6c76193164426ab080783af1c3175b4b999b", size = 11533131, upload-time = "2026-07-12T20:21:35.189Z" },
+    { url = "https://files.pythonhosted.org/packages/5f/fc/04eec35e05a10e0fea1c6503a290ccc3935efda9c845aff64e83282c1af7/ty-0.0.59-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:043c2e00eb1d7475f928af7dedd71f69b64e69bfca55e36f4c968479e1373fc4", size = 12205932, upload-time = "2026-07-12T20:21:37.324Z" },
+    { url = "https://files.pythonhosted.org/packages/a9/dd/a61de859659fa11b55917ad38340a8f2c61f5ae17d1874929f29084c6990/ty-0.0.59-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:f0d688d857441df57f48fca66c029d85cf737c510e7be1d01144cdad1e58d968", size = 12758406, upload-time = "2026-07-12T20:21:39.525Z" },
+    { url = "https://files.pythonhosted.org/packages/c6/e8/fa66f05997eab8ca75fc4f17320140e25467849e0cc75597f898cc22099c/ty-0.0.59-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a96c9f88394a3b42c737e2125b2330543f0d90a43b49761f377d96f8c3ee0d62", size = 12288176, upload-time = "2026-07-12T20:21:41.784Z" },
+    { url = "https://files.pythonhosted.org/packages/15/68/0fca59963bd5123f42d5f7da50667e7a52e8e9615e3a16d8c2c0d3b2d143/ty-0.0.59-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f08dbcb268edcafcb152e59475b5b495ce28d0b340a395c09943557678f4d5a6", size = 12028471, upload-time = "2026-07-12T20:21:43.82Z" },
+    { url = "https://files.pythonhosted.org/packages/e1/5b/cd7dabbbab392578f11179919da5c25d8c3322e5388a688f539ea0539603/ty-0.0.59-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:8812764b9a40fdc98df1272826e73a298ef56b06681135e643bcf90aad1896f7", size = 12297646, upload-time = "2026-07-12T20:21:45.76Z" },
+    { url = "https://files.pythonhosted.org/packages/1d/37/2e9c94f0b383d8cbe1a35517ab470b7810bc9d7501603ab532bcd5be5e90/ty-0.0.59-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:fd53b8581641d8dad7bfac6d5ea589e91a883d6837e0b9a286fdae30722b7c69", size = 11432519, upload-time = "2026-07-12T20:21:47.694Z" },
+    { url = "https://files.pythonhosted.org/packages/9d/0a/af93e9785200f11ac416cc20235fc2464c9bd978e791190684ea0e458795/ty-0.0.59-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:86da5872124a41877d95058bc17d33ddcff034b587eb5f1e2917ab88ba227dac", size = 11554993, upload-time = "2026-07-12T20:21:49.671Z" },
+    { url = "https://files.pythonhosted.org/packages/4b/dd/651bf87e20d00376c81b19124756491cffaf20eb8bec05a8794e5a8cf641/ty-0.0.59-py3-none-musllinux_1_2_i686.whl", hash = "sha256:6a233eef5f2fd4d894881e4a0aec83c9f172bfae1d787d6596ee1939fcc7723e", size = 11818230, upload-time = "2026-07-12T20:21:51.659Z" },
+    { url = "https://files.pythonhosted.org/packages/16/50/c947c4155fea751d135b19affdf734bbce72a94e446b866cf0c62f8bed69/ty-0.0.59-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:7ff678c18b5f1e3128b75a35e50dee7908dea55155baa31cd790619d5014cbf5", size = 12135194, upload-time = "2026-07-12T20:21:53.796Z" },
+    { url = "https://files.pythonhosted.org/packages/b1/13/e5feb138888de1e95037c843571bbbd4ac21bf0a190507468098599a321f/ty-0.0.59-py3-none-win32.whl", hash = "sha256:cf8abb4b8095c5fe39102b8127f5886db308c8d4600909ddbc905512ce9c8163", size = 11179249, upload-time = "2026-07-12T20:21:55.752Z" },
+    { url = "https://files.pythonhosted.org/packages/76/dd/52914dcbeeba92c207de40ef7109a58dcb5527aeb21c8f8feb7402aa9e29/ty-0.0.59-py3-none-win_amd64.whl", hash = "sha256:1dde20a82243d24407869e5a608c2f15efddd5cefc662aef461a5af84bfb3f8b", size = 12251079, upload-time = "2026-07-12T20:21:58.1Z" },
+    { url = "https://files.pythonhosted.org/packages/d4/8f/ac36fde77e223297454c1e0aeb8888c169eaacf3163bb609e3af942c88cb/ty-0.0.59-py3-none-win_arm64.whl", hash = "sha256:987043ee9e021f49493d9135891ac69c1affeee0d4ad4480c5fa4d9c975fc91b", size = 11650921, upload-time = "2026-07-12T20:22:00.348Z" },
 ]
 
 [[package]]

From ca3763ab200a07517f70adca8571a0605e7b6fb3 Mon Sep 17 00:00:00 2001
From: Brad Hilton <brad.hilton.nw@gmail.com>
Date: Wed, 15 Jul 2026 01:54:54 +0000
Subject: [PATCH 2/3] Tighten ty migration typing

---
 src/art/backend.py                            |  24 ++-
 src/art/local/backend.py                      |  18 ++-
 src/art/loss.py                               |   2 +-
 src/art/megatron/dsv4/lora.py                 |   2 +-
 src/art/megatron/dsv4/tokenizer.py            |   2 +-
 .../megatron/flex_attn/flash_dlse_patch.py    |   6 +-
 src/art/megatron/lora.py                      |  46 ++++--
 .../megatron/model_support/handlers/dsv4.py   |   2 +-
 .../megatron/model_support/handlers/gemma4.py |   2 +-
 .../model_support/handlers/gpt_oss.py         |   2 +-
 .../model_support/handlers/qwen3_common.py    |   2 +-
 src/art/megatron/runtime/bridge_runtime.py    |   4 +-
 src/art/megatron/train.py                     |  19 ++-
 src/art/megatron/training/microbatches.py     |  52 ++++---
 src/art/pipeline_trainer/trainer.py           |  19 ++-
 src/art/preprocessing/pack.py                 |   3 +-
 src/art/preprocessing/response_masking.py     |  15 +-
 src/art/preprocessing/tokenize.py             |  46 ++++--
 src/art/serverless/backend.py                 |  10 +-
 src/art/tinker_native/backend.py              |   3 +-
 src/art/trainer_rank/__init__.py              |  13 +-
 src/art/trajectories.py                       |   4 +-
 src/art/unsloth/train.py                      |   2 +-
 .../megatron_attention_oracle_worker.py       |  16 +-
 .../megatron/lora/test_dynamic_lora_slots.py  |  16 +-
 .../megatron/model_support/oracle_worker.py   |  18 ++-
 tests/unit/test_dedicated_config.py           |   2 +-
 tests/unit/test_frontend_logging.py           |   2 +-
 tests/unit/test_moe_routing_real_path.py      |   2 +-
 .../test_pipeline_trainer_local_backend.py    |  16 +-
 tests/unit/test_preprocessing_tokenize.py     |  16 +-
 ...test_serverless_pipeline_trainer_compat.py |   8 +-
 tests/unit/test_tinker_renderers.py           |   9 +-
 tests/unit/test_track_api_cost.py             |   9 +-
 tests/unit/test_trainer_rank_validation.py    | 141 ++++++++++--------
 tests/unit/test_trainer_rank_weird_shapes.py  |  30 ++--
 36 files changed, 356 insertions(+), 227 deletions(-)

diff --git a/src/art/backend.py b/src/art/backend.py
index ed95992c2..617bb6c88 100644
--- a/src/art/backend.py
+++ b/src/art/backend.py
@@ -1,7 +1,16 @@
-from typing import TYPE_CHECKING, Any, AsyncIterator, Iterable, Protocol, TypeAlias
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncIterator,
+    Callable,
+    Coroutine,
+    Iterable,
+    Protocol,
+    TypeAlias,
+)
 
 from . import dev
-from .trajectories import Trajectory, TrajectoryGroup
+from .trajectories import Trajectory
 from .types import TrainResult, TrainSFTConfig
 
 if TYPE_CHECKING:
@@ -35,12 +44,11 @@ async def _prepare_backend_for_training(
         config: dev.OpenAIServerConfig | None,
     ) -> tuple[str, str]: ...
 
-    async def train(
-        self,
-        model: AnyTrainableModel,
-        trajectory_groups: Iterable[TrajectoryGroup],
-        **kwargs: Any,
-    ) -> TrainResult: ...
+    # Backends intentionally expose backend-specific optional training arguments.
+    # Callable[..., ...] preserves that extensibility without falsely requiring
+    # every implementation to accept every other backend's keyword arguments.
+    @property
+    def train(self) -> Callable[..., Coroutine[Any, Any, TrainResult]]: ...
 
     def _train_sft(
         self,
diff --git a/src/art/local/backend.py b/src/art/local/backend.py
index 9c8a9842e..27b5b6f48 100644
--- a/src/art/local/backend.py
+++ b/src/art/local/backend.py
@@ -54,7 +54,7 @@
     aggregate_rl_training_metrics,
     build_rl_train_configs,
 )
-from ..backend import AnyTrainableModel, Backend
+from ..backend import AnyTrainableModel
 from ..dev.sequence_lengths import max_seq_length_from_model_config
 from ..metrics_taxonomy import (
     TRAIN_GRADIENT_STEPS_KEY,
@@ -197,7 +197,7 @@ def _load_training_tokenizer(base_model: str) -> PreTrainedTokenizerBase:
     )
 
 
-class LocalBackend(Backend):
+class LocalBackend:
     def __init__(
         self,
         *,
@@ -834,7 +834,7 @@ async def train(  # type: ignore[override]
         save_checkpoint: bool = True,
         # Verbosity
         verbose: bool = False,
-    ) -> LocalTrainResult:  # ty:ignore[invalid-method-override]
+    ) -> LocalTrainResult:
         """Train the model on the given trajectory groups.
 
         This method does NOT automatically log trajectories or metrics. Call
@@ -1076,10 +1076,11 @@ async def _train_model(
                 try:
                     # Register the copied checkpoint as a new LoRA adapter
                     # so it's available for inference at the new step
-                    if hasattr(service, "register_lora_for_step"):
-                        await service.register_lora_for_step(  # type: ignore[attr-defined]
-                            next_step, next_checkpoint_dir
-                        )  # ty:ignore[call-non-callable]
+                    register_lora_for_step = getattr(
+                        service, "register_lora_for_step", None
+                    )
+                    if callable(register_lora_for_step):
+                        await register_lora_for_step(next_step, next_checkpoint_dir)
                     logger.info(
                         f"[BACKEND] _train_model SKIP: register_lora_for_step "
                         f"completed for step {next_step}"
@@ -1372,7 +1373,8 @@ async def _experimental_pull_model_checkpoint(
                     f"No checkpoints found for {model.project}/{model.name} in local storage or S3"
                 )
             elif local_latest_step is None:
-                resolved_step = s3_latest_step  # type: ignore[assignment]  # ty:ignore[invalid-assignment]
+                assert s3_latest_step is not None
+                resolved_step = s3_latest_step
                 if verbose:
                     print(f"Using latest checkpoint from S3: step {resolved_step}")
             elif s3_latest_step is None:
diff --git a/src/art/loss.py b/src/art/loss.py
index 870b9377d..5fb9269c3 100644
--- a/src/art/loss.py
+++ b/src/art/loss.py
@@ -73,7 +73,7 @@ def align_inputs(self) -> AlignedLossInputs:
             weights=shift_tensor(inputs["weights"], 0.0),
             group_ids=shift_tensor(inputs["group_ids"], 0),
             original_logprobs=(
-                shift_tensor(inputs["original_logprobs"], 0.0)  # ty: ignore[invalid-key, invalid-argument-type]
+                shift_tensor(inputs["original_logprobs"], 0.0)
                 if "original_logprobs" in inputs
                 else None
             ),
diff --git a/src/art/megatron/dsv4/lora.py b/src/art/megatron/dsv4/lora.py
index 7b52c38f6..944b3b769 100644
--- a/src/art/megatron/dsv4/lora.py
+++ b/src/art/megatron/dsv4/lora.py
@@ -619,7 +619,7 @@ def static_config_run(*args: Any, **kwargs: Any) -> Any:
             cache[key] = _dsv4_te_permutation_config(int(key[0]))
         return original_run(*args, **kwargs)
 
-    static_config_run._art_dsv4_static_config_wrapped = True  # type: ignore[attr-defined]  # ty:ignore[unresolved-attribute]
+    setattr(static_config_run, "_art_dsv4_static_config_wrapped", True)
     autotuner.run = static_config_run
     autotuner._art_dsv4_static_config_wrapped = True
 
diff --git a/src/art/megatron/dsv4/tokenizer.py b/src/art/megatron/dsv4/tokenizer.py
index e9eed3f90..03d44b643 100644
--- a/src/art/megatron/dsv4/tokenizer.py
+++ b/src/art/megatron/dsv4/tokenizer.py
@@ -28,7 +28,7 @@ def get_dsv4_tokenizer(
     added_vocab_size = len(added_vocab)
     tokenizer_vocab_size = tokenizer.vocab_size
 
-    class _ArtDsv4Tokenizer(tokenizer.__class__):  # type: ignore[misc, valid-type]  # ty:ignore[unsupported-base]
+    class _ArtDsv4Tokenizer(tokenizer.__class__):  # type: ignore
         def apply_chat_template(
             self,
             messages: list[dict[str, Any]],
diff --git a/src/art/megatron/flex_attn/flash_dlse_patch.py b/src/art/megatron/flex_attn/flash_dlse_patch.py
index 29eee3203..f80bab880 100644
--- a/src/art/megatron/flex_attn/flash_dlse_patch.py
+++ b/src/art/megatron/flex_attn/flash_dlse_patch.py
@@ -138,8 +138,10 @@ def bwd_postprocess_convert_art(
             cluster_size=cluster_size,
         )
 
-    bwd_postprocess_convert_art.compile_cache = (  # type: ignore[attr-defined]  # ty:ignore[unresolved-attribute]
-        original_bwd_postprocess_convert.compile_cache
+    setattr(
+        bwd_postprocess_convert_art,
+        "compile_cache",
+        original_bwd_postprocess_convert.compile_cache,
     )
     cute_interface_any._bwd_postprocess_convert = bwd_postprocess_convert_art
     _DQ_POSTPROCESS_PATCH_APPLIED = True
diff --git a/src/art/megatron/lora.py b/src/art/megatron/lora.py
index 295c2afc9..47f8f3c1b 100644
--- a/src/art/megatron/lora.py
+++ b/src/art/megatron/lora.py
@@ -874,8 +874,12 @@ def _collect_templates(
                 for suffix, param in module._lora_params(slot_ref):
                     if not module._should_export_parameter(param):
                         continue
-                    sharded = bool(param.lora_tp_sharded)  # type: ignore[attr-defined]  # ty:ignore[unresolved-attribute]
-                    shard_domain = param.lora_shard_domain  # type: ignore[attr-defined]  # ty:ignore[unresolved-attribute]
+                    sharded = bool(getattr(param, "lora_tp_sharded"))
+                    shard_domain = getattr(param, "lora_shard_domain")
+                    if shard_domain not in ("tp", "expert_tp"):
+                        raise RuntimeError(
+                            f"invalid LoRA shard domain: {shard_domain!r}"
+                        )
                     templates.append(
                         _LoraPublishTemplate(
                             adapter_model_prefix=module.adapter_model_prefix,
@@ -1046,6 +1050,13 @@ def _expert_grouped_lora_forward(
     return lora(x, tokens_per_expert=tokens_per_expert)
 
 
+def _out_features(module: object) -> int:
+    out_features = getattr(module, "out_features", None)
+    if not isinstance(out_features, int):
+        raise TypeError(f"{type(module).__name__} has no integer out_features")
+    return out_features
+
+
 @torch.compiler.disable
 def _expert_grouped_lora_dual_forward(
     module: "MLPExpertsLinearFC1LoRA",
@@ -1056,9 +1067,7 @@ def _expert_grouped_lora_dual_forward(
     if isinstance(counts, list):
         counts = torch.tensor(counts, dtype=torch.int64, device="cpu")
     if x.shape[0] == 0:
-        return x.new_zeros(
-            (x.shape[0], module.linear_fc1.out_features)  # ty: ignore[unresolved-attribute]
-        )
+        return x.new_zeros((x.shape[0], module.out_features))
     gate = module.gate_lora.active_lora_tensors()
     up = module.up_lora.active_lora_tensors()
     if gate is None or up is None:
@@ -1475,12 +1484,13 @@ def __init__(
     ) -> None:
         super().__init__()
         self.linear_fc1 = linear_fc1
+        self.out_features = _out_features(linear_fc1)
         self.fused_gate_up = bool(fused_gate_up)
         if self.fused_gate_up:
             self.lora = _parallel_lora(
                 adapter_model_prefix=f"{adapter_model_prefix}.{{expert}}.gate_up_proj",
                 linear=linear_fc1,
-                out_features=linear_fc1.out_features,  # ty: ignore[unresolved-attribute]
+                out_features=self.out_features,
                 rank=rank,
                 alpha=alpha,
                 layout="column",
@@ -1489,7 +1499,7 @@ def __init__(
                 allreduce=False,
                 num_local_experts=num_local_experts,
             )
-            gate_out_features = linear_fc1.out_features // 2  # ty: ignore[unresolved-attribute]
+            gate_out_features = self.out_features // 2
             expert_tp_world_size = _get_shard_world_size("expert_tp")
             _set_lora_shard_strategy_metadata(
                 self.lora.B_T,
@@ -1503,7 +1513,7 @@ def __init__(
             self.gate_lora, self.up_lora = _parallel_lora_pair(
                 adapter_model_prefix=f"{adapter_model_prefix}.{{expert}}",
                 linear=linear_fc1,
-                out_features=linear_fc1.out_features // 2,  # ty: ignore[unresolved-attribute]
+                out_features=self.out_features // 2,
                 rank=rank,
                 alpha=alpha,
                 layout="column",
@@ -1526,7 +1536,7 @@ def forward(
                 self.lora,
                 x,
                 tokens_per_expert,
-                self.linear_fc1.out_features,  # ty: ignore[unresolved-attribute]
+                self.out_features,
             )
             if self.fused_gate_up
             else _expert_grouped_lora_dual_forward(self, x, tokens_per_expert)
@@ -1545,10 +1555,11 @@ def __init__(
     ) -> None:
         super().__init__()
         self.linear_fc2 = linear_fc2
+        self.out_features = _out_features(linear_fc2)
         self.lora = _parallel_lora(
             adapter_model_prefix=f"{adapter_model_prefix}.{{expert}}.down_proj",
             linear=linear_fc2,
-            out_features=linear_fc2.out_features,  # ty: ignore[unresolved-attribute]
+            out_features=self.out_features,
             rank=rank,
             alpha=alpha,
             layout="row",
@@ -1572,7 +1583,7 @@ def forward(
             self.lora,
             x,
             tokens_per_expert,
-            self.linear_fc2.out_features,  # ty: ignore[unresolved-attribute]
+            self.out_features,
         )
         # the reason there is no TP comm here is because the MoE token routing handles
         # expert TP comm externally
@@ -1711,7 +1722,7 @@ def wrap_standard_self_attention(
             "linear_qkv",
             TELayerNormColumnParallelLinear,
         )
-        self_attention.linear_qkv = SelfAttentionLinearQKVLoRA(  # ty: ignore[invalid-assignment]
+        linear_qkv_lora = SelfAttentionLinearQKVLoRA(
             adapter_model_prefix=f"{adapter_model_prefix}.self_attn",
             linear_qkv=self_attention_linear_qkv,
             rank=rank,
@@ -1719,6 +1730,7 @@ def wrap_standard_self_attention(
             provider=provider,
             target_modules=target_modules,
         )
+        setattr(self_attention, "linear_qkv", linear_qkv_lora)
 
 
 def wrap_gated_delta_net_attention(
@@ -1776,9 +1788,9 @@ def wrap_grouped_moe_experts(
         mlp_experts_linear_fc1 = _unwrap_attr(
             experts.linear_fc1,
             "linear_fc1",
-            TEColumnParallelGroupedLinear,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+            TEColumnParallelGroupedLinear,  # type: ignore
         )
-        experts.linear_fc1 = MLPExpertsLinearFC1LoRA(  # ty: ignore[invalid-assignment]
+        linear_fc1_lora = MLPExpertsLinearFC1LoRA(
             adapter_model_prefix=f"{adapter_model_prefix}.mlp.experts",
             linear_fc1=mlp_experts_linear_fc1,
             rank=rank,
@@ -1786,6 +1798,7 @@ def wrap_grouped_moe_experts(
             num_local_experts=experts.num_local_experts,
             fused_gate_up=fused_gate_up,
         )
+        setattr(experts, "linear_fc1", linear_fc1_lora)
     wrap_fc2 = (
         wrap_fc1 if fused_gate_up else _targets_include(target_modules, "down_proj")
     )
@@ -1793,15 +1806,16 @@ def wrap_grouped_moe_experts(
         linear_fc2 = _unwrap_attr(
             experts.linear_fc2,
             "linear_fc2",
-            TERowParallelGroupedLinear,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+            TERowParallelGroupedLinear,  # type: ignore
         )
-        experts.linear_fc2 = MLPExpertsLinearFC2LoRA(  # ty: ignore[invalid-assignment]
+        linear_fc2_lora = MLPExpertsLinearFC2LoRA(
             adapter_model_prefix=f"{adapter_model_prefix}.mlp.experts",
             linear_fc2=linear_fc2,
             rank=rank,
             alpha=alpha,
             num_local_experts=experts.num_local_experts,
         )
+        setattr(experts, "linear_fc2", linear_fc2_lora)
 
 
 def wrap_split_mlp_lora(
diff --git a/src/art/megatron/model_support/handlers/dsv4.py b/src/art/megatron/model_support/handlers/dsv4.py
index 213683c15..83d78ecf6 100644
--- a/src/art/megatron/model_support/handlers/dsv4.py
+++ b/src/art/megatron/model_support/handlers/dsv4.py
@@ -268,7 +268,7 @@ def preprocess_hook(
                     )
                 return tuple(preproc_output)
 
-            gpt_module._preprocess = preprocess_hook  # type: ignore[attr-defined]  # ty:ignore[invalid-assignment]
+            setattr(gpt_module, "_preprocess", preprocess_hook)
 
     def collect_layer_families(self, provider: Any) -> list[LayerFamilyInstance]:
         ratios: list[int] = list(getattr(provider, "dsv4_compress_ratios", ()) or ())
diff --git a/src/art/megatron/model_support/handlers/gemma4.py b/src/art/megatron/model_support/handlers/gemma4.py
index 65a09cbbe..8a37c8842 100644
--- a/src/art/megatron/model_support/handlers/gemma4.py
+++ b/src/art/megatron/model_support/handlers/gemma4.py
@@ -799,7 +799,7 @@ def preprocess_hook(
             )
             return tuple(preproc_output)
 
-        gpt_module._preprocess = preprocess_hook  # type: ignore[attr-defined]  # ty:ignore[invalid-assignment]
+        setattr(gpt_module, "_preprocess", preprocess_hook)
 
 
 def _install_gemma4_full_recompute_patch(model_chunks: Sequence[Any]) -> None:
diff --git a/src/art/megatron/model_support/handlers/gpt_oss.py b/src/art/megatron/model_support/handlers/gpt_oss.py
index c4f8e5b6a..c59f08fa6 100644
--- a/src/art/megatron/model_support/handlers/gpt_oss.py
+++ b/src/art/megatron/model_support/handlers/gpt_oss.py
@@ -480,7 +480,7 @@ def preprocess_hook(
             )
             return tuple(preproc_output)
 
-        gpt_module._preprocess = preprocess_hook  # type: ignore[attr-defined]  # ty:ignore[invalid-assignment]
+        setattr(gpt_module, "_preprocess", preprocess_hook)
 
 
 def _install_weighted_bias_quick_geglu_patch() -> None:
diff --git a/src/art/megatron/model_support/handlers/qwen3_common.py b/src/art/megatron/model_support/handlers/qwen3_common.py
index 74ba22b51..0b0c56820 100644
--- a/src/art/megatron/model_support/handlers/qwen3_common.py
+++ b/src/art/megatron/model_support/handlers/qwen3_common.py
@@ -140,4 +140,4 @@ def preprocess_hook(*args, _preprocess=preprocess, **kwargs):
             )
             return tuple(preproc_output)
 
-        gpt_module._preprocess = preprocess_hook  # type: ignore[attr-defined]  # ty:ignore[invalid-assignment]
+        setattr(gpt_module, "_preprocess", preprocess_hook)
diff --git a/src/art/megatron/runtime/bridge_runtime.py b/src/art/megatron/runtime/bridge_runtime.py
index 269087b96..5782aafbf 100644
--- a/src/art/megatron/runtime/bridge_runtime.py
+++ b/src/art/megatron/runtime/bridge_runtime.py
@@ -658,13 +658,13 @@ def _replicated_hf_to_megatron(
         broadcast_device = _materialization_device()
     if self.tp_rank == 0:
         tensor = hf_weights.to(
-            device=cast(Any, broadcast_device),  # ty: ignore[redundant-cast]
+            device=broadcast_device,
             non_blocking=True,
         )
     else:
         tensor = torch.empty_like(
             hf_weights,
-            device=cast(Any, broadcast_device),  # ty: ignore[redundant-cast]
+            device=broadcast_device,
         )
     return self.broadcast_tensor_to_tp_ranks(tensor, src_rank=0)
 
diff --git a/src/art/megatron/train.py b/src/art/megatron/train.py
index 7de21a9c9..c7caef1d9 100644
--- a/src/art/megatron/train.py
+++ b/src/art/megatron/train.py
@@ -919,14 +919,23 @@ def load_adapter_into_model(
     with torch.no_grad():
         for chunk in model_chunks:
             for module in chunk.modules():
-                if hasattr(module, "load_lora"):
-                    module.load_lora(adapter_model)  # type: ignore[attr-defined]  # ty:ignore[call-non-callable]
+                load_lora = getattr(module, "load_lora", None)
+                if callable(load_lora):
+                    load_lora(adapter_model)
 
     if optimizer is None:
         return
     optimizer.reload_model_params()
 
 
+def _zero_grad_buffers(model_chunks: ModelChunks) -> None:
+    for chunk in model_chunks:
+        zero_grad_buffer = getattr(chunk, "zero_grad_buffer", None)
+        if not callable(zero_grad_buffer):
+            raise TypeError(f"{type(chunk).__name__} has no zero_grad_buffer method")
+        zero_grad_buffer()
+
+
 def _optimizer_step(
     optimizer: Any,
     learning_rate: float,
@@ -1314,8 +1323,7 @@ def run_megatron_sft_step(
         moe_routing_replay_controller,
     )
 
-    for chunk in model_chunks:
-        chunk.zero_grad_buffer()  # type: ignore[call-non-callable]  # ty:ignore[call-non-callable]
+    _zero_grad_buffers(model_chunks)
 
     raw_loss_sum: torch.Tensor | None = None
     loss_inputs_for_count: list[dict[str, torch.Tensor] | PreparedSFTMicroInputs] = []
@@ -1469,8 +1477,7 @@ def run_training_step(
     if cp_lookahead_state is not None and int(topology.cp) <= 1:
         cp_lookahead_state.pending_prepared_micro = None
 
-    for chunk in model_chunks:
-        chunk.zero_grad_buffer()  # type: ignore[call-non-callable]  # ty:ignore[call-non-callable]
+    _zero_grad_buffers(model_chunks)
 
     micro_count = len(micro_inputs)
     raw_loss_sum: torch.Tensor | None = None
diff --git a/src/art/megatron/training/microbatches.py b/src/art/megatron/training/microbatches.py
index a80d5529f..36446685a 100644
--- a/src/art/megatron/training/microbatches.py
+++ b/src/art/megatron/training/microbatches.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from typing import Any
 
 from megatron.core import parallel_state as ps
@@ -56,6 +56,24 @@ class PreparedSFTMicroInputs(BaseModel):
     local_token_uids: torch.Tensor | None = None
 
 
+def _map_packed_tensors(
+    inputs: PackedTensors,
+    transform: Callable[[torch.Tensor], torch.Tensor],
+) -> PackedTensors:
+    mapped = inputs.copy()
+    mapped["tokens"] = transform(inputs["tokens"])
+    mapped["group_ids"] = transform(inputs["group_ids"])
+    mapped["parent_ids"] = transform(inputs["parent_ids"])
+    mapped["input_pos"] = transform(inputs["input_pos"])
+    mapped["assistant_mask"] = transform(inputs["assistant_mask"])
+    mapped["logprobs"] = transform(inputs["logprobs"])
+    mapped["advantages"] = transform(inputs["advantages"])
+    mapped["weights"] = transform(inputs["weights"])
+    if "original_logprobs" in inputs:
+        mapped["original_logprobs"] = transform(inputs["original_logprobs"])
+    return mapped
+
+
 @torch.no_grad()
 def select_indexed_inputs(packed_tensors: PackedTensors, index: int) -> PackedTensors:
     def selected_tensor(value: torch.Tensor) -> torch.Tensor:
@@ -65,28 +83,20 @@ def selected_tensor(value: torch.Tensor) -> torch.Tensor:
             return selected.clone()
         return selected
 
-    return PackedTensors(  # type: ignore[call-arg]  # ty:ignore[missing-typed-dict-key]
-        **{
-            key: selected_tensor(value)
-            for key, value in packed_tensors.items()
-            if isinstance(value, torch.Tensor)
-        },
-        pixel_values=[None],
-        image_grid_thw=[None],
-    )
+    selected = _map_packed_tensors(packed_tensors, selected_tensor)
+    selected.pop("moe_routing_replay", None)
+    selected["pixel_values"] = [None]
+    selected["image_grid_thw"] = [None]
+    return selected
 
 
 @torch.no_grad()
 def _clone_packed_tensors(inputs: PackedTensors) -> PackedTensors:
-    return PackedTensors(  # type: ignore[call-arg]  # ty:ignore[missing-typed-dict-key]
-        **{
-            key: value.clone()
-            for key, value in inputs.items()
-            if isinstance(value, torch.Tensor)
-        },
-        pixel_values=[None],
-        image_grid_thw=[None],
-    )
+    cloned = _map_packed_tensors(inputs, torch.Tensor.clone)
+    cloned.pop("moe_routing_replay", None)
+    cloned["pixel_values"] = [None]
+    cloned["image_grid_thw"] = [None]
+    return cloned
 
 
 @torch.no_grad()
@@ -225,9 +235,7 @@ def _select_next_step_first_micro(
 
 
 def _move_inputs_to_device(inputs: PackedTensors, device: torch.device) -> None:
-    for key, value in inputs.items():
-        if isinstance(value, torch.Tensor):
-            inputs[key] = value.to(device)  # type: ignore[index]  # ty:ignore[invalid-key]
+    inputs.update(_map_packed_tensors(inputs, lambda tensor: tensor.to(device)))
 
 
 def _count_trainable_tokens(inputs: LossInputs | DispatchedPackedTensors) -> float:
diff --git a/src/art/pipeline_trainer/trainer.py b/src/art/pipeline_trainer/trainer.py
index 8a9d214b8..f280bb34d 100644
--- a/src/art/pipeline_trainer/trainer.py
+++ b/src/art/pipeline_trainer/trainer.py
@@ -20,6 +20,8 @@
     cast,
 )
 
+from typing_extensions import TypeIs
+
 T = TypeVar("T")
 
 import art
@@ -42,9 +44,17 @@
 _ACTOR_IDLE_TIME_KEY = "_art_actor_idle_s"
 
 
+def _is_eval_mapping(
+    result: Sequence[art.Trajectory | art.TrajectoryGroup]
+    | Mapping[str, Sequence[art.Trajectory | art.TrajectoryGroup]],
+) -> TypeIs[Mapping[str, Sequence[art.Trajectory | art.TrajectoryGroup]]]:
+    return isinstance(result, Mapping)
+
+
 def _to_async_iterator(iterable: Iterable[T] | AsyncIterator[T]) -> AsyncIterator[T]:
     """Convert a sync Iterable to an AsyncIterator, or pass through if already async."""
     if isinstance(iterable, AsyncIterator):
+        # ty cannot currently preserve T through this runtime generic check.
         return cast(AsyncIterator[T], iterable)
 
     async def _iter():
@@ -769,14 +779,7 @@ async def _run_eval(self, step: int) -> None:
             finally:
                 token.var.reset(token)
                 eval_elapsed = time.monotonic() - eval_started
-            splits: Mapping[str, Sequence[art.Trajectory | art.TrajectoryGroup]]
-            if isinstance(result, Mapping):
-                splits = cast(
-                    Mapping[str, Sequence[art.Trajectory | art.TrajectoryGroup]],
-                    result,
-                )
-            else:
-                splits = {"val": result}
+            splits = result if _is_eval_mapping(result) else {"val": result}
 
             logged_eval_timing = False
             for split_name, items in splits.items():
diff --git a/src/art/preprocessing/pack.py b/src/art/preprocessing/pack.py
index e42558947..215655b5f 100644
--- a/src/art/preprocessing/pack.py
+++ b/src/art/preprocessing/pack.py
@@ -34,7 +34,8 @@ class PackedTensors(TypedDict):
     weights: torch.Tensor
     pixel_values: list[torch.Tensor | None]
     image_grid_thw: list[torch.Tensor | None]
-    moe_routing_replay: PackedMoeRoutingReplay | None
+    moe_routing_replay: NotRequired[PackedMoeRoutingReplay | None]
+    original_logprobs: NotRequired[torch.Tensor]
 
 
 class DiskPackedTensors(TypedDict):
diff --git a/src/art/preprocessing/response_masking.py b/src/art/preprocessing/response_masking.py
index 8ee31efef..24a243ca5 100644
--- a/src/art/preprocessing/response_masking.py
+++ b/src/art/preprocessing/response_masking.py
@@ -1,11 +1,20 @@
-from transformers.tokenization_utils_base import PreTrainedTokenizerBase
+from collections.abc import Iterable
+from typing import Protocol
+
+
+class _TemplatePartTokenizer(Protocol):
+    def __call__(self, text: str, *, add_special_tokens: bool) -> object: ...
 
 
 def token_ids_for_template_part(
-    tokenizer: PreTrainedTokenizerBase,
+    tokenizer: _TemplatePartTokenizer,
     template_part: str,
 ) -> list[int]:
-    return list(tokenizer(template_part, add_special_tokens=False).input_ids)
+    encoded = tokenizer(template_part, add_special_tokens=False)
+    input_ids = getattr(encoded, "input_ids", None)
+    if not isinstance(input_ids, Iterable):
+        raise TypeError("tokenizer must return iterable input_ids")
+    return [int(token_id) for token_id in input_ids]
 
 
 def _find_subsequence(
diff --git a/src/art/preprocessing/tokenize.py b/src/art/preprocessing/tokenize.py
index dc866f3d4..9d8ca0c34 100644
--- a/src/art/preprocessing/tokenize.py
+++ b/src/art/preprocessing/tokenize.py
@@ -7,7 +7,7 @@
 import json
 import math
 import random
-from typing import TYPE_CHECKING, Any, Generator, Literal, cast
+from typing import TYPE_CHECKING, Any, Generator, Literal, Protocol, cast
 
 from openai.types.chat.chat_completion import Choice
 import torch
@@ -27,15 +27,31 @@
     TokenRoute,
     align_choice_routes_to_tokenized_result,
 )
-from .response_masking import response_only_labels, token_ids_for_template_part
+from .response_masking import (
+    _TemplatePartTokenizer,
+    response_only_labels,
+    token_ids_for_template_part,
+)
 from .vllm_tokens import choice_vllm_token_metadata
 
 ChatTemplateTool = dict[Any, Any] | Callable[..., Any]
 ChatTemplateToolSchemaFormat = Literal["default", "vllm_openai"]
 
 
+class _TokenDecoder(Protocol):
+    def decode(self, token_ids: int, /) -> str | list[str]: ...
+
+
+class _SFTTokenizer(_TemplatePartTokenizer, Protocol):
+    @property
+    def chat_template(self) -> object: ...
+
+    @property
+    def apply_chat_template(self) -> Callable[..., object]: ...
+
+
 def _chat_template_kwargs(
-    tokenizer: PreTrainedTokenizerBase,
+    tokenizer: _SFTTokenizer,
     chat_template_kwargs: dict[str, Any] | None,
 ) -> dict[str, Any]:
     return merge_chat_template_kwargs(
@@ -93,7 +109,7 @@ def _normalize_tools_for_chat_template(
 
 
 def _normalize_tool_call_arguments_for_chat_template(
-    tokenizer: PreTrainedTokenizerBase,
+    tokenizer: _SFTTokenizer,
     messages: list[dict[str, Any]],
 ) -> list[dict[str, Any]]:
     chat_template = tokenizer.chat_template
@@ -127,7 +143,7 @@ def _normalize_tool_call_arguments_for_chat_template(
 
 
 def _messages_for_chat_template(
-    tokenizer: PreTrainedTokenizerBase,
+    tokenizer: _SFTTokenizer,
     messages_and_choices: MessagesAndChoices,
     *,
     final_trainable_choice_index: int | None = None,
@@ -159,7 +175,7 @@ class TokenizedResult:
     trajectory: Trajectory
     choice_offsets: list[int]
     extra_logprobs: dict[str, list[float]]
-    _tokenizer: "PreTrainedTokenizerBase" = field(repr=False, compare=False)
+    _tokenizer: _TokenDecoder = field(repr=False, compare=False)
     moe_routed_experts: list[TokenRoute | None] | None = None
     moe_routing_alignment_stats: MoeRoutingAlignmentStats | None = None
     weight: float = 0.0
@@ -168,9 +184,13 @@ class TokenizedResult:
 
     @cached_property
     def tokens(self) -> list[str]:
-        return [
-            cast(str, self._tokenizer.decode(token_id)) for token_id in self.token_ids
-        ]
+        tokens: list[str] = []
+        for token_id in self.token_ids:
+            token = self._tokenizer.decode(token_id)
+            if not isinstance(token, str):
+                raise TypeError("decoding one token must return one string")
+            tokens.append(token)
+        return tokens
 
     def without_prompt(self) -> "TokenizedResult":
         return TokenizedResult(
@@ -230,7 +250,7 @@ def _validate_max_seq_length(max_seq_length: int | None) -> None:
 
 
 def _apply_chat_template_token_ids(
-    tokenizer: PreTrainedTokenizerBase,
+    tokenizer: _SFTTokenizer,
     messages: list[dict[str, Any]],
     **kwargs: Any,
 ) -> list[int]:
@@ -247,7 +267,7 @@ def _apply_chat_template_token_ids(
 
 
 def _apply_chat_template_text(
-    tokenizer: PreTrainedTokenizerBase,
+    tokenizer: _SFTTokenizer,
     messages: list[dict[str, Any]],
     **kwargs: Any,
 ) -> str:
@@ -262,7 +282,7 @@ def _apply_chat_template_text(
 
 
 def _last_assistant_input_ids_and_labels(
-    tokenizer: PreTrainedTokenizerBase,
+    tokenizer: _SFTTokenizer,
     messages: list[dict[str, Any]],
     tools: list[ChatTemplateTool] | None,
     template_kwargs: dict[str, Any],
@@ -592,7 +612,7 @@ def tokenize_trajectory(
 def tokenize_sft_batch(
     trajectory_batch: list[Trajectory],
     learning_rate: float,
-    tokenizer: PreTrainedTokenizerBase,
+    tokenizer: _SFTTokenizer,
     instruction_part: str,
     response_part: str,
     chat_template_kwargs: dict[str, Any] | None = None,
diff --git a/src/art/serverless/backend.py b/src/art/serverless/backend.py
index ea83475d1..44c4dd279 100644
--- a/src/art/serverless/backend.py
+++ b/src/art/serverless/backend.py
@@ -15,7 +15,7 @@
     aggregate_rl_training_metrics,
     build_rl_train_configs,
 )
-from ..backend import AnyTrainableModel, Backend
+from ..backend import AnyTrainableModel
 from ..metrics_taxonomy import (
     TRAIN_GRADIENT_STEPS_KEY,
     build_training_summary_metrics,
@@ -86,7 +86,7 @@ def _canonicalize_upstream_metrics(metrics: dict[str, float]) -> dict[str, float
     }
 
 
-class ServerlessBackend(Backend):
+class ServerlessBackend:
     def __init__(
         self, *, api_key: str | None = None, base_url: str | None = None
     ) -> None:
@@ -257,7 +257,7 @@ async def train(  # type: ignore[override]
         save_checkpoint: bool = True,
         # Verbosity
         verbose: bool = False,
-    ) -> ServerlessTrainResult:  # ty:ignore[invalid-method-override]
+    ) -> ServerlessTrainResult:
         """Train the model on the given trajectory groups.
 
         This method does NOT automatically log trajectories or metrics. Call
@@ -816,7 +816,7 @@ async def _experimental_pull_from_s3(
 
     async def _experimental_push_to_s3(
         self,
-        model: "Model",
+        model: "TrainableModel",
         *,
         s3_bucket: str | None = None,
         prefix: str | None = None,
@@ -858,7 +858,7 @@ async def _experimental_push_to_s3(
 
             # Pull from W&B to local temp dir
             checkpoint_dir = await self._experimental_pull_model_checkpoint(
-                model,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+                model,
                 step=step,
                 verbose=verbose,
             )
diff --git a/src/art/tinker_native/backend.py b/src/art/tinker_native/backend.py
index 1701c6cc4..0ae9093b1 100644
--- a/src/art/tinker_native/backend.py
+++ b/src/art/tinker_native/backend.py
@@ -29,7 +29,6 @@
 import uvicorn
 
 from .. import dev
-from ..backend import Backend
 from ..costs import build_cost_calculator, compute_train_cost, get_model_pricing
 from ..metrics_taxonomy import (
     build_training_summary_metrics,
@@ -178,7 +177,7 @@ class TinkerNativeModelConfig:
     training_client_args: dict[str, Any]
 
 
-class TinkerNativeBackend(Backend):
+class TinkerNativeBackend:
     _tinker_train_log_env = "ART_TINKER_TRAIN_LOG"
     _tinker_sample_log_env = "ART_TINKER_SAMPLE_LOG"
 
diff --git a/src/art/trainer_rank/__init__.py b/src/art/trainer_rank/__init__.py
index c60bb3a13..3d10f87b4 100644
--- a/src/art/trainer_rank/__init__.py
+++ b/src/art/trainer_rank/__init__.py
@@ -22,6 +22,7 @@
 
 import torch
 import torch.distributed as dist
+from typing_extensions import TypeIs
 
 from art.megatron.prefix_tree_packing import (
     PrefixTreePack,
@@ -2453,7 +2454,7 @@ def _gather_tensor_parallel_logits(self, logits: torch.Tensor) -> torch.Tensor:
         )
 
 
-def _validate_top_k(top_k: int, model: "GPTModel") -> None:
+def _validate_top_k(top_k: int, model: object) -> None:
     vocab_size = _padded_vocab_size(model)
     if top_k > vocab_size:
         raise ValueError(f"top_k={top_k} exceeds vocabulary size {vocab_size}")
@@ -2519,7 +2520,7 @@ def _language_model(model: torch.nn.Module) -> "GPTModel":
     raise RuntimeError("expected a Megatron GPT model")
 
 
-def _padded_vocab_size(model: "GPTModel") -> int:
+def _padded_vocab_size(model: object) -> int:
     vocab_size = getattr(getattr(model, "config", None), "padded_vocab_size", None)
     if vocab_size is None:
         vocab_size = getattr(model, "vocab_size", None)
@@ -2874,9 +2875,13 @@ def _materialize(inputs: ForwardInputs) -> ForwardInputs:
     return [_materialize(item) for item in _nested_forward_children(inputs)]
 
 
+def _is_forward_input(inputs: ForwardInputs) -> TypeIs[AnyForwardInput]:
+    return isinstance(inputs, ForwardInput)
+
+
 def _flatten(inputs: ForwardInputs) -> Iterator[AnyForwardInput]:
-    if isinstance(inputs, ForwardInput):
-        yield inputs  # ty: ignore[invalid-yield]
+    if _is_forward_input(inputs):
+        yield inputs
         return
     for item in _nested_forward_children(inputs):
         yield from _flatten(item)
diff --git a/src/art/trajectories.py b/src/art/trajectories.py
index b844ac41f..a929d0e71 100644
--- a/src/art/trajectories.py
+++ b/src/art/trajectories.py
@@ -123,7 +123,7 @@ def get_messages(messages_and_choices: MessagesAndChoices) -> Messages:
             msg = dict(message_or_choice)
             if msg.get("content") is None:
                 msg["content"] = ""
-            messages.append(msg)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+            messages.append(cast(Message, msg))
     return messages
 
 
@@ -223,7 +223,7 @@ def __deepcopy__(self, memo: dict[int, Any] | None = None):
     def log(self, message: str) -> None:
         self.logs.append(message)
 
-    def __iter__(self) -> Iterator[Trajectory]:  # type: ignore[override]  # ty:ignore[invalid-method-override]
+    def __iter__(self) -> Iterator[Trajectory]:  # type: ignore
         return iter(self.trajectories)
 
     def __len__(self) -> int:
diff --git a/src/art/unsloth/train.py b/src/art/unsloth/train.py
index ca95d7d67..91c73711f 100644
--- a/src/art/unsloth/train.py
+++ b/src/art/unsloth/train.py
@@ -830,7 +830,7 @@ async def run_unsloth_rl_training(
     for offset in range(0, packed_tensors["tokens"].shape[0]):
         for _ in range(2 if warmup else 1):
             if precalculate_logprobs and not warmup:
-                packed_tensors["original_logprobs"] = packed_tensors["logprobs"]  # type: ignore[index]  # ty:ignore[invalid-key]
+                packed_tensors["original_logprobs"] = packed_tensors["logprobs"]
                 packed_tensors["logprobs"] = _precalculate_new_logprobs(
                     ctx,
                     packed_tensors,
diff --git a/tests/integration/megatron/cp_attn/megatron_attention_oracle_worker.py b/tests/integration/megatron/cp_attn/megatron_attention_oracle_worker.py
index 0bdca2e15..a98ff58b3 100644
--- a/tests/integration/megatron/cp_attn/megatron_attention_oracle_worker.py
+++ b/tests/integration/megatron/cp_attn/megatron_attention_oracle_worker.py
@@ -69,12 +69,14 @@ def run_worker_subprocess(
         "--run-request",
         str(request_path),
     ]
-    env = {
-        **os.environ,
-        "ART_MEGATRON_ATTACH_TOKEN_UIDS": "1",
-        "PYTHONUNBUFFERED": "1",
-        "PYTHONPATH": os.pathsep.join(pythonpath_entries),
-    }
+    env = dict(os.environ)
+    env.update(
+        {
+            "ART_MEGATRON_ATTACH_TOKEN_UIDS": "1",
+            "PYTHONUNBUFFERED": "1",
+            "PYTHONPATH": os.pathsep.join(pythonpath_entries),
+        }
+    )
     env.pop("ART_FLEX_BACKEND", None)
     for cache_env in ("TORCHINDUCTOR_CACHE_DIR", "TRITON_CACHE_DIR"):
         cache_root = env.get(cache_env)
@@ -97,7 +99,7 @@ def run_worker_subprocess(
         log_file.flush()
         live_log_file.write(header)
         live_log_file.flush()
-        run = subprocess.Popen(  # ty: ignore[no-matching-overload]
+        run: subprocess.Popen[bytes] = subprocess.Popen(
             command,
             cwd=str(worker_cwd),
             env=env,
diff --git a/tests/integration/megatron/lora/test_dynamic_lora_slots.py b/tests/integration/megatron/lora/test_dynamic_lora_slots.py
index c8d3335f5..bc6ddd817 100644
--- a/tests/integration/megatron/lora/test_dynamic_lora_slots.py
+++ b/tests/integration/megatron/lora/test_dynamic_lora_slots.py
@@ -62,7 +62,9 @@ def test_dynamic_lora_slots_capture_recompute_context_and_step_independently() -
         assert torch.allclose(actual, expected, atol=0, rtol=0)
         assert slot_a.rank == 1
         assert slot_a.scale == 32.0
-        assert lora._slot(ref_b).scale == 8.0  # type: ignore[union-attr]  # ty:ignore[unresolved-attribute]
+        slot_b = lora._slot(ref_b)
+        assert slot_b is not None
+        assert slot_b.scale == 8.0
 
         trainer = _trainer_for(lora, device)
         cpu_adapter = {
@@ -326,8 +328,12 @@ def _assert_checkpoint_recomputes_with(
         y = checkpoint(lambda t: lora(t), *checkpoint_args, x)
     with use_lora_slot(ambient_ref):
         y.sum().backward()
-    assert lora._slot(expected_ref).A_T.grad is not None  # type: ignore[union-attr]  # ty:ignore[unresolved-attribute]
-    assert lora._slot(ambient_ref).A_T.grad is None  # type: ignore[union-attr]  # ty:ignore[unresolved-attribute]
+    expected_slot = lora._slot(expected_ref)
+    ambient_slot = lora._slot(ambient_ref)
+    assert expected_slot is not None
+    assert ambient_slot is not None
+    assert expected_slot.A_T.grad is not None
+    assert ambient_slot.A_T.grad is None
 
 
 def _assert_step_updates_only(
@@ -375,7 +381,9 @@ def _assert_reload_replaces_slot_optimizer(
     assert ref.name not in trainer._dynamic_optimizers
     assert [tuple(param.shape) for param in new_params] == [(4, 3), (3, 5)]
     assert all(old is not new for old, new in zip(old_params, new_params, strict=True))
-    assert lora._slot(ref).rank == 3  # type: ignore[union-attr]  # ty:ignore[unresolved-attribute]
+    slot = lora._slot(ref)
+    assert slot is not None
+    assert slot.rank == 3
 
 
 def _trainer_for(lora: LoRA, device: torch.device) -> TrainerRank:
diff --git a/tests/integration/megatron/model_support/oracle_worker.py b/tests/integration/megatron/model_support/oracle_worker.py
index a845bd370..98a75acc1 100644
--- a/tests/integration/megatron/model_support/oracle_worker.py
+++ b/tests/integration/megatron/model_support/oracle_worker.py
@@ -525,7 +525,7 @@ def _apply_requested_flex_backend_patch(flex_backend: str | None):
     else:
         raise RuntimeError(f"Unsupported flex backend request: {flex_backend}")
 
-    compiled_flex_attention._FORCED_FLEX_BACKEND = patched_backend  # type: ignore[invalid-assignment]  # ty:ignore[invalid-assignment]
+    setattr(compiled_flex_attention, "_FORCED_FLEX_BACKEND", patched_backend)
     compiled_flex_attention._FORCED_FLEX_KERNEL_OPTIONS = patched_kernel_options
     compiled_flex_attention.dense_compiled_flex_attention = torch.compile(
         compiled_flex_attention._forced_flex_attention_dense
@@ -697,9 +697,11 @@ def _assert_runtime_configuration(
     standard_attention_layers = 0
 
     try:
-        from megatron.core.ssm.gated_delta_net import GatedDeltaNet
+        import megatron.core.ssm.gated_delta_net as gated_delta_net
     except ImportError:  # pragma: no cover - optional dependency guard.
-        GatedDeltaNet = ()  # type: ignore[assignment]  # ty:ignore[invalid-assignment]
+        gated_delta_net_type = None
+    else:
+        gated_delta_net_type = getattr(gated_delta_net, "GatedDeltaNet")
     from megatron.core.transformer.attention import SelfAttention
 
     for chunk in model_chunks:
@@ -712,7 +714,9 @@ def _assert_runtime_configuration(
         if config is not None and hasattr(config, "context_parallel_size"):
             observed_context_parallel_sizes.add(int(config.context_parallel_size))
         for child in module.modules():
-            if GatedDeltaNet and isinstance(child, GatedDeltaNet):
+            if gated_delta_net_type is not None and isinstance(
+                child, gated_delta_net_type
+            ):
                 gdn_layers += 1
             if isinstance(child, SelfAttention):
                 standard_attention_layers += 1
@@ -1044,7 +1048,7 @@ def _mutated_sanitize(grad: torch.Tensor | None) -> torch.Tensor | None:
         view.copy_(grad)
         return view
 
-    executor._sanitize_nested_stage_input_grad = _mutated_sanitize  # type: ignore[invalid-assignment]  # ty:ignore[invalid-assignment]
+    setattr(executor, "_sanitize_nested_stage_input_grad", _mutated_sanitize)
     try:
         yield
     finally:
@@ -1066,8 +1070,8 @@ def _apply_attention_lse_normalize_mutation(mutation: SensitivityMutation | None
     def _identity(lse: torch.Tensor, **_kwargs: Any) -> torch.Tensor:
         return lse
 
-    compiled_flex_attention.normalize_flex_lse = _identity  # type: ignore[invalid-assignment]  # ty:ignore[invalid-assignment]
-    executor.normalize_flex_lse = _identity  # type: ignore[invalid-assignment]  # ty:ignore[invalid-assignment]
+    setattr(compiled_flex_attention, "normalize_flex_lse", _identity)
+    setattr(executor, "normalize_flex_lse", _identity)
     try:
         yield
     finally:
diff --git a/tests/unit/test_dedicated_config.py b/tests/unit/test_dedicated_config.py
index be5d7c583..4f2fe9308 100644
--- a/tests/unit/test_dedicated_config.py
+++ b/tests/unit/test_dedicated_config.py
@@ -248,7 +248,7 @@ def test_invalid_rollout_weights_mode():
         ValueError, match="rollout_weights_mode must be either 'lora' or 'merged'"
     ):
         validate_dedicated_config(
-            InternalModelConfig(rollout_weights_mode="bad-mode")  # type: ignore[typeddict-item]  # ty:ignore[invalid-argument-type]
+            InternalModelConfig(rollout_weights_mode="bad-mode")  # type: ignore
         )
 
 
diff --git a/tests/unit/test_frontend_logging.py b/tests/unit/test_frontend_logging.py
index 31a12ae2b..6a622e305 100644
--- a/tests/unit/test_frontend_logging.py
+++ b/tests/unit/test_frontend_logging.py
@@ -1264,7 +1264,7 @@ async def mock_train_model(*args, **kwargs):
                     TRAIN_GRADIENT_STEPS_KEY: 2.0,
                 }
 
-        backend._train_model = mock_train_model  # type: ignore[method-assign]  # ty:ignore[invalid-assignment]
+        setattr(backend, "_train_model", mock_train_model)
         backend._get_step = AsyncMock(return_value=1)  # type: ignore[method-assign]
 
         groups = [
diff --git a/tests/unit/test_moe_routing_real_path.py b/tests/unit/test_moe_routing_real_path.py
index 244e6ef2b..369df614d 100644
--- a/tests/unit/test_moe_routing_real_path.py
+++ b/tests/unit/test_moe_routing_real_path.py
@@ -141,7 +141,7 @@ def _tokenized(
         trajectory=Trajectory(),
         choice_offsets=[trainable_start],
         extra_logprobs={},
-        _tokenizer=_FakeTokenizer(),  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+        _tokenizer=_FakeTokenizer(),
         moe_routed_experts=cast(list[list[list[int]] | None], routes),
         prompt_id=prompt_id,
         prompt_length=prompt_length,
diff --git a/tests/unit/test_pipeline_trainer_local_backend.py b/tests/unit/test_pipeline_trainer_local_backend.py
index 32e57e0d5..554fd0603 100644
--- a/tests/unit/test_pipeline_trainer_local_backend.py
+++ b/tests/unit/test_pipeline_trainer_local_backend.py
@@ -55,7 +55,7 @@ def _make_trainer(
 ) -> PipelineTrainer:
     return PipelineTrainer(
         model=model,
-        backend=backend,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+        backend=backend,  # type: ignore
         rollout_fn=_noop_rollout,
         scenarios=[],
         config={},
@@ -239,7 +239,8 @@ async def test_pipeline_trainer_uses_same_train_kwargs_for_local_backend(
         ),
     )
     backend = LocalBackend(path=str(tmp_path))
-    backend.train = AsyncMock(return_value=SimpleNamespace(step=1, metrics={}))  # type: ignore[method-assign]
+    train = AsyncMock(return_value=SimpleNamespace(step=1, metrics={}))
+    setattr(backend, "train", train)
 
     trainer = _make_trainer(
         model=model,
@@ -253,7 +254,8 @@ async def test_pipeline_trainer_uses_same_train_kwargs_for_local_backend(
 
     await trainer._training_stage()
 
-    assert backend.train.await_args.kwargs == {  # type: ignore[attr-defined]  # ty:ignore[unresolved-attribute]
+    assert train.await_args is not None
+    assert train.await_args.kwargs == {
         "learning_rate": 3e-5,
         "loss_fn": "ppo",
         "loss_fn_config": None,
@@ -286,7 +288,7 @@ async def fake_train_model(
         seen["verbose"] = verbose
         yield {}
 
-    backend._train_model = fake_train_model  # type: ignore[method-assign]  # ty:ignore[invalid-assignment]
+    setattr(backend, "_train_model", fake_train_model)
     backend._get_step = AsyncMock(return_value=1)  # type: ignore[method-assign]
     with patch.object(model, "_get_wandb_run", return_value=None):
         result = await backend.train(
@@ -326,7 +328,7 @@ async def fake_train_model(
         seen["verbose"] = verbose
         yield {}
 
-    backend._train_model = fake_train_model  # type: ignore[method-assign]  # ty:ignore[invalid-assignment]
+    setattr(backend, "_train_model", fake_train_model)
     backend._get_step = AsyncMock(return_value=1)  # type: ignore[method-assign]
     with patch.object(model, "_get_wandb_run", return_value=None):
         result = await backend.train(
@@ -367,7 +369,7 @@ async def fake_train_model(
         seen["dev_config"] = dev_config
         yield {}
 
-    backend._train_model = fake_train_model  # type: ignore[method-assign]  # ty:ignore[invalid-assignment]
+    setattr(backend, "_train_model", fake_train_model)
     backend._get_step = AsyncMock(return_value=1)  # type: ignore[method-assign]
 
     with patch.object(model, "_get_wandb_run", return_value=None):
@@ -409,7 +411,7 @@ async def fake_train_model(
         seen["dev_config"] = dev_config
         yield {}
 
-    backend._train_model = fake_train_model  # type: ignore[method-assign]  # ty:ignore[invalid-assignment]
+    setattr(backend, "_train_model", fake_train_model)
     backend._get_step = AsyncMock(return_value=1)  # type: ignore[method-assign]
     with patch.object(model, "_get_wandb_run", return_value=None):
         await backend.train(
diff --git a/tests/unit/test_preprocessing_tokenize.py b/tests/unit/test_preprocessing_tokenize.py
index 46c6afc4c..ebf9f1f8d 100644
--- a/tests/unit/test_preprocessing_tokenize.py
+++ b/tests/unit/test_preprocessing_tokenize.py
@@ -181,7 +181,7 @@ def test_tokenize_sft_batch_masks_response_tokens_without_unsloth_import() -> No
     batch = tokenize_sft_batch(
         trajectory_batch=[Trajectory(messages_and_choices=messages, reward=1.0)],
         learning_rate=1e-5,
-        tokenizer=tokenizer,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+        tokenizer=tokenizer,
         instruction_part="<user>",
         response_part="<assistant>",
     )
@@ -207,7 +207,7 @@ def test_tokenize_sft_batch_all_trains_every_assistant_turn() -> None:
     batch = tokenize_sft_batch(
         trajectory_batch=[Trajectory(messages_and_choices=messages)],
         learning_rate=1e-5,
-        tokenizer=tokenizer,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+        tokenizer=tokenizer,
         instruction_part="<user>",
         response_part="<assistant>",
         assistant_turns="all",
@@ -231,7 +231,7 @@ def test_tokenize_sft_batch_passes_chat_template_kwargs() -> None:
     tokenize_sft_batch(
         trajectory_batch=[Trajectory(messages_and_choices=messages, reward=1.0)],
         learning_rate=1e-5,
-        tokenizer=tokenizer,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+        tokenizer=tokenizer,
         instruction_part="<user>",
         response_part="<assistant>",
         chat_template_kwargs={
@@ -269,7 +269,7 @@ def test_tokenize_sft_batch_trains_only_final_assistant_turn() -> None:
     batch = tokenize_sft_batch(
         trajectory_batch=[Trajectory(messages_and_choices=messages)],
         learning_rate=1e-5,
-        tokenizer=tokenizer,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+        tokenizer=tokenizer,
         instruction_part="<user>",
         response_part="<assistant>",
         assistant_turns="last",
@@ -313,7 +313,7 @@ def test_tokenize_sft_batch_trains_final_tool_call_and_turn_end() -> None:
     batch = tokenize_sft_batch(
         trajectory_batch=[Trajectory(messages_and_choices=messages)],
         learning_rate=1e-5,
-        tokenizer=tokenizer,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+        tokenizer=tokenizer,
         instruction_part="<user>",
         response_part="<assistant>",
         assistant_turns="last",
@@ -358,7 +358,7 @@ def test_tokenize_sft_batch_preserves_prompt_target_token_boundary() -> None:
     batch = tokenize_sft_batch(
         trajectory_batch=[Trajectory(messages_and_choices=messages)],
         learning_rate=1e-5,
-        tokenizer=tokenizer,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+        tokenizer=tokenizer,
         instruction_part="<user>",
         response_part="<assistant>",
         assistant_turns="last",
@@ -380,7 +380,7 @@ def test_tokenize_sft_batch_last_requires_final_assistant_message() -> None:
                 Trajectory(messages_and_choices=[{"role": "user", "content": "hello"}])
             ],
             learning_rate=1e-5,
-            tokenizer=tokenizer,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+            tokenizer=tokenizer,
             instruction_part="<user>",
             response_part="<assistant>",
             assistant_turns="last",
@@ -400,7 +400,7 @@ def test_tokenize_sft_batch_last_rejects_non_prefix_stable_template() -> None:
                 )
             ],
             learning_rate=1e-5,
-            tokenizer=tokenizer,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+            tokenizer=tokenizer,
             instruction_part="<user>",
             response_part="<assistant>",
             assistant_turns="last",
diff --git a/tests/unit/test_serverless_pipeline_trainer_compat.py b/tests/unit/test_serverless_pipeline_trainer_compat.py
index 7a083fd21..bf914c49c 100644
--- a/tests/unit/test_serverless_pipeline_trainer_compat.py
+++ b/tests/unit/test_serverless_pipeline_trainer_compat.py
@@ -56,7 +56,7 @@ async def fake_train_model(
         seen["verbose"] = verbose
         yield {"loss": 0.25}
 
-    backend._train_model = fake_train_model  # type: ignore[method-assign]  # ty:ignore[invalid-assignment]
+    setattr(backend, "_train_model", fake_train_model)
     backend._get_step = AsyncMock(return_value=3)  # type: ignore[method-assign]
 
     with patch.object(model, "_get_wandb_run", return_value=None):
@@ -156,7 +156,7 @@ async def test_serverless_train_model_forwards_experimental_config() -> None:
     async def events_list(**_kwargs: Any):
         yield SimpleNamespace(id="event-id", type="training_ended", data={})
 
-    backend._client.training_jobs.events.list = events_list  # type: ignore[attr-defined]  # ty:ignore[invalid-assignment]
+    setattr(backend._client.training_jobs.events, "list", events_list)  # type: ignore[attr-defined]
 
     async def no_sleep(_seconds: float) -> None:
         return None
@@ -227,7 +227,7 @@ async def test_serverless_train_sft_forwards_metric_logging_config() -> None:
     async def events_list(**_kwargs: Any):
         yield SimpleNamespace(id="event-id", type="training_ended", data={})
 
-    backend._client.sft_training_jobs.events.list = events_list  # type: ignore[attr-defined]  # ty:ignore[invalid-assignment]
+    setattr(backend._client.sft_training_jobs.events, "list", events_list)  # type: ignore[attr-defined]
 
     async def no_sleep(_seconds: float) -> None:
         return None
@@ -332,7 +332,7 @@ async def fake_train_model(
         seen["dev_config"] = dev_config
         yield {}
 
-    backend._train_model = fake_train_model  # type: ignore[method-assign]  # ty:ignore[invalid-assignment]
+    setattr(backend, "_train_model", fake_train_model)
     backend._get_step = AsyncMock(return_value=1)  # type: ignore[method-assign]
 
     with patch.object(model, "_get_wandb_run", return_value=None):
diff --git a/tests/unit/test_tinker_renderers.py b/tests/unit/test_tinker_renderers.py
index 87f3d4422..46ef9cfff 100644
--- a/tests/unit/test_tinker_renderers.py
+++ b/tests/unit/test_tinker_renderers.py
@@ -1,6 +1,7 @@
 import json
 from typing import cast
 
+from tinker import EncodedTextChunk, ModelInput
 from tinker_cookbook import renderers
 from tinker_cookbook.tokenizer_utils import Tokenizer
 
@@ -51,10 +52,12 @@ def decode(self, tokens: int | list[int]) -> str:
         return "".join(self._id_to_text[token] for token in tokens)
 
 
-def _decode_model_input(tokenizer: FakeTokenizer, model_input: object) -> str:
+def _decode_model_input(tokenizer: FakeTokenizer, model_input: ModelInput) -> str:
     tokens: list[int] = []
-    for chunk in model_input.chunks:  # type: ignore[attr-defined]  # ty:ignore[unresolved-attribute]
-        assert hasattr(chunk, "tokens"), f"Unexpected non-text chunk: {chunk!r}"
+    for chunk in model_input.chunks:
+        assert isinstance(chunk, EncodedTextChunk), (
+            f"Unexpected non-text chunk: {chunk!r}"
+        )
         tokens.extend(list(chunk.tokens))
     return tokenizer.decode(tokens)
 
diff --git a/tests/unit/test_track_api_cost.py b/tests/unit/test_track_api_cost.py
index dc4969241..9a5ad92ed 100644
--- a/tests/unit/test_track_api_cost.py
+++ b/tests/unit/test_track_api_cost.py
@@ -712,10 +712,17 @@ async def eval_fn(
                 )
             ]
 
+        async def rollout_fn(
+            _model: TrainableModel,
+            _scenario: dict,
+            _config: dict,
+        ) -> TrajectoryGroup:
+            return TrajectoryGroup([])
+
         trainer = PipelineTrainer(
             model=model,
             backend=backend,
-            rollout_fn=lambda *_args, **_kwargs: asyncio.sleep(0),  # ty: ignore[invalid-argument-type]
+            rollout_fn=rollout_fn,
             scenarios=[],
             config={},
             num_rollout_workers=1,
diff --git a/tests/unit/test_trainer_rank_validation.py b/tests/unit/test_trainer_rank_validation.py
index 8dfbe88b5..bdd21ecb6 100644
--- a/tests/unit/test_trainer_rank_validation.py
+++ b/tests/unit/test_trainer_rank_validation.py
@@ -1,16 +1,18 @@
 from __future__ import annotations
 
+from collections.abc import Iterable
 from dataclasses import dataclass
 import gc
 import inspect
 from types import SimpleNamespace
-from typing import Any, cast
+from typing import TYPE_CHECKING, Any, cast
 
 import pytest
 import torch
 
 from art.trainer_rank import (
     AdamParams,
+    AdapterSelection,
     ForwardInput,
     ForwardOutput,
     TopK,
@@ -24,6 +26,10 @@
     _validate_top_k,
 )
 
+if TYPE_CHECKING:
+    from art.megatron.lora import LoRASlotRef
+    from art.megatron.train import TrainingRuntime
+
 
 class _Model:
     vocab_size = 8
@@ -72,13 +78,19 @@ def _runtime(
     model: torch.nn.Module | None = None,
     *,
     optimizer: object | None = None,
-) -> SimpleNamespace:
+) -> "TrainingRuntime":
+    # Deliberately lightweight structural fake; importing/constructing the real
+    # Megatron runtime would make these CPU-only unit tests require Megatron.
     return SimpleNamespace(
         model=[model or torch.nn.Linear(1, 1)],
         optimizer=optimizer,
         provider=SimpleNamespace(hidden_size=4, num_layers=1),
         model_support_handler=SimpleNamespace(build_gdn_execution_spec=True),
-    )
+    )  # type: ignore
+
+
+def _slot_ref(kind: str, name: str | None) -> "LoRASlotRef":
+    return _SlotRef(kind, name)  # type: ignore
 
 
 def _target_request(token: int) -> ForwardInput[torch.Tensor, None, None, None]:
@@ -110,7 +122,8 @@ def _output_values(outputs: object) -> list[int]:
         assert isinstance(target_logprobs, torch.Tensor)
         return [int(target_logprobs.item())]
     values: list[int] = []
-    for item in outputs:  # type: ignore[union-attr]  # ty:ignore[not-iterable]
+    assert isinstance(outputs, Iterable)
+    for item in outputs:
         values.extend(_output_values(item))
     return values
 
@@ -118,14 +131,15 @@ def _output_values(outputs: object) -> list[int]:
 def _output_shape(outputs: object) -> object:
     if isinstance(outputs, ForwardOutput):
         return "output"
-    return [_output_shape(item) for item in outputs]  # type: ignore[union-attr]  # ty:ignore[not-iterable]
+    assert isinstance(outputs, Iterable)
+    return [_output_shape(item) for item in outputs]
 
 
 def _trainer_with_checkpoint(
     monkeypatch: pytest.MonkeyPatch,
     value: torch.Tensor,
 ) -> tuple[TrainerRank, torch.nn.Parameter]:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime())
     param = torch.nn.Parameter(value.clone())
     trainer._checkpoint_slot_params_by_name["student"] = (param,)
     monkeypatch.setattr(
@@ -137,16 +151,21 @@ def _trainer_with_checkpoint(
 
 
 def _tracked_targets(
-    trainer: TrainerRank, ref: _SlotRef, *scales: float
+    trainer: TrainerRank, ref: "LoRASlotRef", *scales: float
 ) -> list[torch.Tensor]:
     tracked = trainer._track_slot_graph_outputs(
-        ref,  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+        ref,
         [
             ForwardOutput(torch.ones(1, requires_grad=True) * scale, None, None, None)
             for scale in scales
         ],
     )
-    return [cast(torch.Tensor, output.target_logprobs) for output in tracked]
+    targets: list[torch.Tensor] = []
+    for output in tracked:
+        target = output.target_logprobs
+        assert isinstance(target, torch.Tensor)
+        targets.append(target)
+    return targets
 
 
 def test_forward_input_validation() -> None:
@@ -155,14 +174,14 @@ def test_forward_input_validation() -> None:
     with pytest.raises(ValueError, match="cannot set both checkpoint and lora"):
         ForwardInput(input_tokens=torch.tensor([1]), checkpoint="a", lora="b")
     with pytest.raises(ValueError, match="top_k=9 exceeds vocabulary size 8"):
-        _validate_top_k(9, _Model())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+        _validate_top_k(9, _Model())
 
 
 @pytest.mark.parametrize(("checkpoint", "expected"), ((Unset, Unset), (None, None)))
 def test_forward_input_distinguishes_unset_and_base_checkpoint(
-    checkpoint: object, expected: object
+    checkpoint: AdapterSelection, expected: AdapterSelection
 ) -> None:
-    request = ForwardInput(input_tokens=torch.tensor([1]), checkpoint=checkpoint)  # type: ignore[arg-type]  # ty:ignore[no-matching-overload]
+    request = ForwardInput(input_tokens=torch.tensor([1]), checkpoint=checkpoint)
 
     assert request.checkpoint is expected
     assert request.lora is Unset
@@ -176,24 +195,24 @@ def test_forward_input_preserves_public_runtime_shape() -> None:
 
 @pytest.mark.parametrize("depth", (0, 2))
 def test_trainer_rank_accepts_shared_prefix_depth(depth: int) -> None:
-    trainer = TrainerRank(_runtime(), shared_prefix_max_depth=depth)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime(), shared_prefix_max_depth=depth)
 
     assert trainer.shared_prefix_max_depth == depth
 
 
 def test_trainer_rank_adapter_stack_errors() -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime())
 
     with pytest.raises(RuntimeError, match="No pushed LoRA or checkpoint"):
         trainer.pop_pushed_lora_or_checkpoint()
-    trainer._slot_stack.append(object())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer._slot_stack.append(object())  # type: ignore
     for load in (trainer.load_checkpoint_slot, trainer.load_lora_slot):
         with pytest.raises(RuntimeError, match="Cannot load a LoRA/checkpoint"):
             load("teacher", {})
 
 
 def test_trainer_rank_rejects_adapter_keys_without_installed_lora_site() -> None:
-    trainer = TrainerRank(_runtime(_FakeLoRASite("base.layer")))  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime(_FakeLoRASite("base.layer")))
     valid = {
         "base.layer.lora_A.weight": torch.empty(1),
         "base.layer.lora_B.weight": torch.empty(1),
@@ -210,7 +229,7 @@ def test_trainer_rank_rejects_adapter_keys_without_installed_lora_site() -> None
 
 def test_trainer_rank_normalizes_adapter_tensors_to_installed_site() -> None:
     site = _FakeLoRASite("base.layer", dtype=torch.bfloat16)
-    trainer = TrainerRank(_runtime(site))  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime(site))
     adapter = {
         "base.layer.lora_A.weight": torch.ones(3, 4, dtype=torch.float32),
         "base.layer.lora_B.weight": torch.ones(5, 3, dtype=torch.float32),
@@ -223,7 +242,7 @@ def test_trainer_rank_normalizes_adapter_tensors_to_installed_site() -> None:
 
 
 def test_checkpoint_slot_adapter_config_is_validated_and_copied() -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime())
     config = {
         "base_model_name_or_path": "Qwen/Qwen3-8B",
         "r": 8,
@@ -250,7 +269,7 @@ def test_checkpoint_slot_adapter_config_is_validated_and_copied() -> None:
 def test_checkpoint_slot_adapter_config_rejects_cross_rank_mismatch(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime())
     monkeypatch.setattr("art.trainer_rank.dist.is_initialized", lambda: True)
     monkeypatch.setattr("art.trainer_rank.dist.get_world_size", lambda: 2)
 
@@ -266,7 +285,7 @@ def gather(output: list[object], value: object) -> None:
 def test_load_checkpoint_slot_retains_config_and_uses_its_alpha(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime())
     seen: dict[str, object] = {}
     monkeypatch.setattr(
         trainer,
@@ -294,7 +313,7 @@ def test_load_checkpoint_slot_retains_config_and_uses_its_alpha(
 
 
 def test_checkpoint_slot_publish_requires_retained_adapter_config() -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime())
     with pytest.raises(ValueError, match="Unknown checkpoint slot"):
         trainer.save_checkpoint_slot_lora("missing", "/unused")
 
@@ -305,7 +324,7 @@ def test_checkpoint_slot_publish_requires_retained_adapter_config() -> None:
 
 
 def test_trainer_rank_default_forward_uses_explicit_base_slot() -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime())
 
     plan = trainer._plan_flat_forward([_target_request(1)])
 
@@ -318,7 +337,7 @@ def test_trainer_rank_default_forward_uses_explicit_base_slot() -> None:
 
 def test_optim_step_requires_loaded_checkpoint_slot() -> None:
     optimizer = _NativeOptimizer()
-    trainer = TrainerRank(_runtime(optimizer=optimizer))  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime(optimizer=optimizer))
 
     with pytest.raises(TrainerRankSlotStateError, match="loaded checkpoint slot"):
         trainer.optim_step(params=AdamParams(learning_rate=1e-3))
@@ -327,7 +346,7 @@ def test_optim_step_requires_loaded_checkpoint_slot() -> None:
 
 
 def test_optim_step_rejects_loaded_slots_without_grads() -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime())
     trainer._checkpoint_slot_params_by_name["student"] = (
         torch.nn.Parameter(torch.ones(2)),
     )
@@ -344,7 +363,7 @@ def test_optim_step_rejects_loaded_slots_without_grads() -> None:
 def test_optim_step_rejects_explicit_slot_subset_with_missing_grads(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime())
     ready = torch.nn.Parameter(torch.ones(2))
     missing = torch.nn.Parameter(torch.ones(2))
     ready.grad = torch.ones_like(ready)
@@ -366,7 +385,7 @@ def test_optim_step_rejects_explicit_slot_subset_with_missing_grads(
 def test_optim_step_implicitly_steps_only_slots_with_grads(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime())
     ready = torch.nn.Parameter(torch.ones(2))
     untouched = torch.nn.Parameter(torch.ones(2))
     ready.grad = torch.ones_like(ready)
@@ -488,12 +507,12 @@ def test_trainer_rank_rejects_mutating_slot_with_pending_graph(
     operation: str,
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
-    ref = _SlotRef("checkpoint", "teacher")
-    monkeypatch.setattr(trainer, "_slot_ref", lambda kind, name: _SlotRef(kind, name))
+    trainer = TrainerRank(_runtime())
+    ref = _slot_ref("checkpoint", "teacher")
+    monkeypatch.setattr(trainer, "_slot_ref", _slot_ref)
     target = _tracked_targets(trainer, ref, 2)[0]
     guard = (
-        (lambda: trainer._guard_slot_can_load(ref))  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+        (lambda: trainer._guard_slot_can_load(ref))
         if operation == "load"
         else (lambda: trainer._guard_checkpoint_can_step("teacher"))
     )
@@ -509,14 +528,14 @@ def test_trainer_rank_step_allows_missing_slot_graph_bookkeeping(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     trainer = TrainerRank.__new__(TrainerRank)
-    monkeypatch.setattr(trainer, "_slot_ref", lambda kind, name: _SlotRef(kind, name))
+    monkeypatch.setattr(trainer, "_slot_ref", _slot_ref)
 
     trainer._guard_checkpoint_can_step("student")
 
 
 def test_trainer_rank_zero_grad_does_not_clear_live_slot_graphs() -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
-    ref = _SlotRef("lora", "teacher")
+    trainer = TrainerRank(_runtime())
+    ref = _slot_ref("lora", "teacher")
     output = ForwardOutput(
         None,
         TopK(
@@ -527,67 +546,67 @@ def test_trainer_rank_zero_grad_does_not_clear_live_slot_graphs() -> None:
         None,
     )
 
-    tracked = trainer._track_slot_graph_outputs(ref, [output])  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    tracked = trainer._track_slot_graph_outputs(ref, [output])
     trainer.zero_grad()
 
     assert tracked[0].top_k is not None
     with pytest.raises(TrainerRankSlotStateError, match="live backward graph"):
-        trainer._guard_slot_can_load(ref)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+        trainer._guard_slot_can_load(ref)
 
 
 def test_trainer_rank_retained_backward_keeps_slot_graph_guard() -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
-    ref = _SlotRef("checkpoint", "teacher")
+    trainer = TrainerRank(_runtime())
+    ref = _slot_ref("checkpoint", "teacher")
     target = _tracked_targets(trainer, ref, 2)[0]
 
     target.sum().backward(retain_graph=True)
     with pytest.raises(TrainerRankSlotStateError, match="live backward graph"):
-        trainer._guard_slot_can_load(ref)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+        trainer._guard_slot_can_load(ref)
 
     target.sum().backward()
-    trainer._guard_slot_can_load(ref)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer._guard_slot_can_load(ref)
 
 
 def test_trainer_rank_tracks_each_independent_output_graph() -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
-    ref = _SlotRef("checkpoint", "teacher")
+    trainer = TrainerRank(_runtime())
+    ref = _slot_ref("checkpoint", "teacher")
     first, second = _tracked_targets(trainer, ref, 2, 3)
 
     first.sum().backward()
     with pytest.raises(TrainerRankSlotStateError, match="live backward graph"):
-        trainer._guard_slot_can_load(ref)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+        trainer._guard_slot_can_load(ref)
 
     second.sum().backward()
-    trainer._guard_slot_can_load(ref)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer._guard_slot_can_load(ref)
 
 
 def test_trainer_rank_tracks_graph_after_output_is_replaced_by_loss() -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
-    ref = _SlotRef("checkpoint", "teacher")
+    trainer = TrainerRank(_runtime())
+    ref = _slot_ref("checkpoint", "teacher")
     target = _tracked_targets(trainer, ref, 2)[0]
     loss = target.sum()
     del target
     gc.collect()
 
     with pytest.raises(TrainerRankSlotStateError, match="live backward graph"):
-        trainer._guard_slot_can_load(ref)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+        trainer._guard_slot_can_load(ref)
 
     loss.backward()
-    trainer._guard_slot_can_load(ref)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer._guard_slot_can_load(ref)
 
 
 def test_trainer_rank_releases_abandoned_output_graph() -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
-    ref = _SlotRef("checkpoint", "teacher")
+    trainer = TrainerRank(_runtime())
+    ref = _slot_ref("checkpoint", "teacher")
     target = _tracked_targets(trainer, ref, 2)[0]
     del target
     gc.collect()
 
-    trainer._guard_slot_can_load(ref)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer._guard_slot_can_load(ref)
 
 
 def test_dp_rank_forward_preserves_nested_shape_for_inactive_requests() -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime())
     request_a = ForwardInput(input_tokens=torch.tensor([1]))
     request_b = ForwardInput(input_tokens=torch.tensor([2]))
 
@@ -604,7 +623,7 @@ def test_dp_rank_forward_preserves_nested_shape_for_inactive_requests() -> None:
 def test_dp_rank_forward_supports_arbitrary_nested_depth(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime())
     _stub_forward(monkeypatch, trainer, _indexed_outputs)
     nested = [
         [[[[[_target_request(1)]]]]],
@@ -623,7 +642,7 @@ def test_dp_rank_forward_supports_arbitrary_nested_depth(
 def test_forward_micro_batches_uses_deterministic_dp_windows(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime())
     _stub_forward(monkeypatch, trainer, dp=(1, 2))
 
     batches = list(
@@ -637,7 +656,7 @@ def test_forward_micro_batches_uses_deterministic_dp_windows(
 def test_forward_micro_batches_syncs_fit_decision_across_dp(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime())
     _stub_forward(monkeypatch, trainer, dp=(1, 2), profiled=True)
     sync_flags: list[bool] = []
 
@@ -659,7 +678,7 @@ def memory_check(required: int, *, sync_across_dp: bool = False) -> _MemoryCheck
 def test_forward_micro_batches_supports_arbitrary_nested_depth(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime())
     _stub_forward(monkeypatch, trainer, _indexed_outputs, profiled=True)
     expected = [
         [[[[[_target_request(1)]]]]],
@@ -680,7 +699,7 @@ def test_forward_micro_batches_supports_arbitrary_nested_depth(
 def test_forward_micro_batches_ramps_after_first_success(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime())
 
     def run(plan, **_kwargs):
         trainer._memory_profiles[plan.signature] = _MemoryProfile(
@@ -706,7 +725,7 @@ def run(plan, **_kwargs):
 def test_forward_micro_batches_does_not_overtrust_tiny_memory_profile(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime())
     monkeypatch.setattr(trainer, "_dp_rank_and_size", lambda: (0, 1))
     inputs = [_target_request(i) for i in range(64)]
     tiny_plan = trainer._plan_flat_forward([inputs[0]])
@@ -724,7 +743,7 @@ def test_forward_micro_batches_does_not_overtrust_tiny_memory_profile(
 def test_forward_micro_batches_tail_does_not_reset_stable_window(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime())
     trainer._last_global_micro_batch_size = 64
     _stub_forward(monkeypatch, trainer, profiled=True)
     monkeypatch.setattr(
@@ -752,7 +771,7 @@ def test_forward_micro_batches_tail_does_not_reset_stable_window(
 def test_forward_micro_batches_raises_when_smallest_batch_will_not_fit(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime())
     monkeypatch.setattr(trainer, "_dp_rank_and_size", lambda: (0, 1))
     monkeypatch.setattr(
         trainer,
@@ -775,7 +794,7 @@ def test_forward_micro_batches_raises_when_smallest_batch_will_not_fit(
 def test_forward_micro_batches_rejects_mismatched_replicated_counts(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    trainer = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime())
     import art.trainer_rank as trainer_rank
 
     monkeypatch.setattr(trainer_rank.dist, "is_available", lambda: True)
@@ -814,7 +833,7 @@ def __init__(self) -> None:
         def _preprocess(self, *args: object, **kwargs: object) -> None:
             return None
 
-    trainer = TrainerRank(_runtime(FakeGPT()))  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    trainer = TrainerRank(_runtime(FakeGPT()))
     tokens = torch.tensor([[1, 2, 3]], dtype=torch.long)
     labels = torch.stack((tokens, tokens + 1), dim=-1)
 
diff --git a/tests/unit/test_trainer_rank_weird_shapes.py b/tests/unit/test_trainer_rank_weird_shapes.py
index a5562db01..62b7ff365 100644
--- a/tests/unit/test_trainer_rank_weird_shapes.py
+++ b/tests/unit/test_trainer_rank_weird_shapes.py
@@ -2,6 +2,7 @@
 
 from collections.abc import Callable, Iterable
 from types import SimpleNamespace
+from typing import TYPE_CHECKING
 
 import pytest
 import torch
@@ -23,6 +24,9 @@
     _MemoryProfile,
 )
 
+if TYPE_CHECKING:
+    from art.megatron.train import TrainingRuntime
+
 
 class _FakeGPT(torch.nn.Module):
     def __init__(self, *, hidden_size: int = 8, vocab_size: int = 32) -> None:
@@ -39,13 +43,15 @@ def _preprocess(self, *args: object, **kwargs: object) -> None:
         return None
 
 
-def _runtime() -> SimpleNamespace:
+def _runtime() -> "TrainingRuntime":
+    # Deliberately lightweight structural fake; importing/constructing the real
+    # Megatron runtime would make these CPU-only unit tests require Megatron.
     return SimpleNamespace(
         model=[_FakeGPT()],
         optimizer=None,
         provider=SimpleNamespace(hidden_size=8, num_layers=4),
         model_support_handler=SimpleNamespace(build_gdn_execution_spec=True),
-    )
+    )  # type: ignore
 
 
 def _tokens(*values: int) -> torch.Tensor:
@@ -208,7 +214,7 @@ def test_shared_trainable_tokens_accumulate_independent_output_gradients() -> No
 
 
 def test_planner_handles_vineppo_nested_shape_and_request_mix() -> None:
-    rank = TrainerRank(_runtime(), shared_prefix_max_depth=3)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    rank = TrainerRank(_runtime(), shared_prefix_max_depth=3)
     inputs = _vineppo_like_inputs()
     flat = list(_flatten(inputs))
 
@@ -231,7 +237,7 @@ def test_planner_handles_vineppo_nested_shape_and_request_mix() -> None:
 def test_forward_micro_batches_preserves_nested_vineppo_groups(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    rank = TrainerRank(_runtime(), shared_prefix_max_depth=2)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    rank = TrainerRank(_runtime(), shared_prefix_max_depth=2)
     monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1))
     monkeypatch.setattr(rank, "_all_ranks_have_memory_profile", lambda **_kwargs: True)
     monkeypatch.setattr(
@@ -264,7 +270,7 @@ def test_forward_micro_batches_preserves_nested_vineppo_groups(
 def test_adaptive_planner_materializes_only_final_large_candidate(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    rank = TrainerRank(_runtime(), shared_prefix_max_depth=3)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    rank = TrainerRank(_runtime(), shared_prefix_max_depth=3)
     rank._last_global_micro_batch_size = 32
     monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1))
     monkeypatch.setattr(rank, "_all_ranks_have_memory_profile", lambda **_kwargs: True)
@@ -310,7 +316,7 @@ def estimate(requests):
 def test_adaptive_planner_globally_falls_back_when_one_rank_cannot_estimate(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    rank = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    rank = TrainerRank(_runtime())
     monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 2))
     monkeypatch.setattr(rank, "_all_ranks_true", lambda _local: False)
     plans = 0
@@ -333,7 +339,7 @@ def plan(requests):
 def test_adaptive_planner_probes_new_heterogeneous_signatures(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    rank = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    rank = TrainerRank(_runtime())
     monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1))
     monkeypatch.setattr(rank, "_resolve_slot_ref", lambda request: request.checkpoint)
     inputs = [
@@ -363,7 +369,7 @@ def test_adaptive_planner_probes_new_heterogeneous_signatures(
 def test_adaptive_planner_grows_stable_window_to_largest_aligned_fit(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    rank = TrainerRank(_runtime(), shared_prefix_max_depth=1)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    rank = TrainerRank(_runtime(), shared_prefix_max_depth=1)
     rank._last_global_micro_batch_size = 512
     monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1))
     monkeypatch.setattr(rank, "_all_ranks_have_memory_profile", lambda **_kwargs: True)
@@ -381,7 +387,7 @@ def test_adaptive_planner_grows_stable_window_to_largest_aligned_fit(
 def test_forward_micro_batches_shrinks_when_memory_budget_drops(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    rank = TrainerRank(_runtime(), shared_prefix_max_depth=2)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    rank = TrainerRank(_runtime(), shared_prefix_max_depth=2)
     monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1))
     monkeypatch.setattr(rank, "_all_ranks_have_memory_profile", lambda **_kwargs: True)
     inputs = [_target_request(_tokens(1, 2, 3, index)) for index in range(14)]
@@ -430,7 +436,7 @@ def run(plan, **_kwargs):
 def test_heterogeneous_slots_split_packing_without_losing_output_estimates(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    rank = TrainerRank(_runtime(), shared_prefix_max_depth=4)  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    rank = TrainerRank(_runtime(), shared_prefix_max_depth=4)
     monkeypatch.setattr(
         TrainerRank,
         "_slot_ref",
@@ -466,7 +472,7 @@ def test_forward_raises_before_expected_oom_with_actionable_context(
     api: str,
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    rank = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    rank = TrainerRank(_runtime())
     if api == "dp_rank_forward":
         monkeypatch.setattr(
             rank,
@@ -508,7 +514,7 @@ def test_flatten_rejects_dicts_to_avoid_silent_top_level_shape_changes() -> None
 
 
 def test_no_output_requests_do_not_pack_or_consume_compute_memory() -> None:
-    rank = TrainerRank(_runtime())  # type: ignore[arg-type]  # ty:ignore[invalid-argument-type]
+    rank = TrainerRank(_runtime())
     requests: Iterable[ForwardInput] = [
         ForwardInput(input_tokens=_tokens(1, 2, 3)),
         ForwardInput(input_tokens=_tokens(1, 2, 4)),

From 5bb9472b5b4033d16036a37000b84c866f303706 Mon Sep 17 00:00:00 2001
From: Brad Hilton <brad.hilton.nw@gmail.com>
Date: Wed, 15 Jul 2026 02:28:16 +0000
Subject: [PATCH 3/3] Fix ty migration runtime regressions

---
 src/art/guided_completion.py              |  6 ++----
 src/art/megatron/training/microbatches.py | 21 ++++++++-------------
 src/art/utils/trajectory_migration.py     |  2 +-
 tests/unit/test_trajectory_parquet.py     | 15 +++++++++++++++
 4 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/src/art/guided_completion.py b/src/art/guided_completion.py
index f652a1d3f..cd1418f2f 100644
--- a/src/art/guided_completion.py
+++ b/src/art/guided_completion.py
@@ -36,7 +36,7 @@ def get_guided_completion_params(
     completion: ChatCompletion,
     base_tools: Iterable[ChatCompletionToolParam] | None = None,
 ) -> Tuple[
-    List[str] | None,
+    List[str | None] | None,
     ChatCompletionToolChoiceOptionParam | None,
     List[ChatCompletionToolParam] | None,
 ]:
@@ -50,7 +50,7 @@ def get_guided_completion_params(
 
     Returns a tuple of (guided_choice, tool_choice, tool_params).
     """
-    guided_choice: List[str] | None = None
+    guided_choice: List[str | None] | None = None
     tool_choice: ChatCompletionToolChoiceOptionParam | None = None
     tool_params: List[ChatCompletionToolParam] | None = None
 
@@ -77,7 +77,5 @@ def get_guided_completion_params(
         ]
     else:
         content = completion.choices[0].message.content
-        if content is None:
-            raise ValueError("Completion has no text or tool call")
         guided_choice = [content]
     return (guided_choice, tool_choice, tool_params)
diff --git a/src/art/megatron/training/microbatches.py b/src/art/megatron/training/microbatches.py
index 36446685a..0f222de61 100644
--- a/src/art/megatron/training/microbatches.py
+++ b/src/art/megatron/training/microbatches.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Callable, Sequence
-from typing import Any
+from typing import Any, cast
 
 from megatron.core import parallel_state as ps
 from pydantic import BaseModel, ConfigDict
@@ -60,18 +60,13 @@ def _map_packed_tensors(
     inputs: PackedTensors,
     transform: Callable[[torch.Tensor], torch.Tensor],
 ) -> PackedTensors:
-    mapped = inputs.copy()
-    mapped["tokens"] = transform(inputs["tokens"])
-    mapped["group_ids"] = transform(inputs["group_ids"])
-    mapped["parent_ids"] = transform(inputs["parent_ids"])
-    mapped["input_pos"] = transform(inputs["input_pos"])
-    mapped["assistant_mask"] = transform(inputs["assistant_mask"])
-    mapped["logprobs"] = transform(inputs["logprobs"])
-    mapped["advantages"] = transform(inputs["advantages"])
-    mapped["weights"] = transform(inputs["weights"])
-    if "original_logprobs" in inputs:
-        mapped["original_logprobs"] = transform(inputs["original_logprobs"])
-    return mapped
+    return cast(
+        PackedTensors,
+        {
+            key: transform(value) if isinstance(value, torch.Tensor) else value
+            for key, value in inputs.items()
+        },
+    )
 
 
 @torch.no_grad()
diff --git a/src/art/utils/trajectory_migration.py b/src/art/utils/trajectory_migration.py
index 7ec947b69..b295ef961 100644
--- a/src/art/utils/trajectory_migration.py
+++ b/src/art/utils/trajectory_migration.py
@@ -79,7 +79,7 @@ def trajectory_to_dict(trajectory: Trajectory) -> dict[str, Any]:
 def message_or_choice_to_dict(message_or_choice: MessageOrChoice) -> dict[str, Any]:
     # messages are sometimes stored as dicts, so we need to handle both cases
     item_dict = (
-        message_or_choice.model_dump()
+        message_or_choice.to_dict()
         if isinstance(message_or_choice, Choice)
         else message_or_choice
     )
diff --git a/tests/unit/test_trajectory_parquet.py b/tests/unit/test_trajectory_parquet.py
index c728bbcb8..05354d499 100644
--- a/tests/unit/test_trajectory_parquet.py
+++ b/tests/unit/test_trajectory_parquet.py
@@ -44,6 +44,7 @@
 from art.utils.trajectory_migration import (
     MigrationResult,
     deserialize_trajectory_groups,
+    message_or_choice_to_dict,
     migrate_jsonl_to_parquet,
     migrate_model_dir,
     migrate_trajectories_dir,
@@ -55,6 +56,20 @@
 FIXTURES_DIR = Path(__file__).parent.parent / "fixtures" / "trajectories"
 
 
+def test_legacy_choice_serialization_omits_unset_fields() -> None:
+    choice = Choice(
+        index=0,
+        finish_reason="stop",
+        message={"role": "assistant", "content": "ok"},
+    )
+
+    assert message_or_choice_to_dict(choice) == {
+        "finish_reason": "stop",
+        "index": 0,
+        "message": {"content": "ok", "role": "assistant"},
+    }
+
+
 def _ensure_message(item: MessageOrChoice) -> ChatCompletionMessageParam:
     """Narrow a trajectory entry to a concrete message (not a Choice)."""
     assert not isinstance(item, Choice)