diff --git a/examples/deepscaler/train_deepscaler_nb.py b/examples/deepscaler/train_deepscaler_nb.py
index e49aea741..4d5d845d7 100644
--- a/examples/deepscaler/train_deepscaler_nb.py
+++ b/examples/deepscaler/train_deepscaler_nb.py
@@ -373,7 +373,7 @@ def process_item(item):
 train_dataset, test_dataset = create_datasets()
 train_dataset, val_dataset = data_lib.post_init_dataset(
     train_dataset,
-    tokenizer,
+    tokenizer,  # pyrefly: ignore[bad-argument-type]
     batch_size=BATCH_SIZE,
     num_batches=NUM_BATCHES,
     max_prompt_length=MAX_PROMPT_LENGTH,
@@ -383,7 +383,7 @@ def process_item(item):
 
 test_dataset, _ = data_lib.post_init_dataset(
     test_dataset,
-    tokenizer,
+    tokenizer,  # pyrefly: ignore[bad-argument-type]
     batch_size=BATCH_SIZE,
     num_batches=NUM_TEST_BATCHES,
     max_prompt_length=MAX_PROMPT_LENGTH,
diff --git a/examples/deepswe/train_deepswe_nb.py b/examples/deepswe/train_deepswe_nb.py
index 761f1d960..2047b5da3 100644
--- a/examples/deepswe/train_deepswe_nb.py
+++ b/examples/deepswe/train_deepswe_nb.py
@@ -645,7 +645,7 @@ def mixed_type_batch_fn(elements):
 
 train_dataset, _ = data_lib.post_init_dataset(
     grain_dataset,
-    tokenizer,
+    tokenizer,  # pyrefly: ignore[bad-argument-type]
     batch_size=BATCH_SIZE,
     num_batches=None,
     max_prompt_length=MAX_PROMPT_LENGTH,
@@ -856,7 +856,7 @@ def get_lora_model(base_model, model_mesh):
     "temperature": TEMPERATURE,
     "top_p": TOP_P,
     "top_k": TOP_K,
-    "eos_tokens": [tokenizer.encode("<|im_end|>")[0]],
+    "eos_tokens": [tokenizer.encode("<|im_end|>")[0]],  # pyrefly: ignore[missing-attribute]
     "return_logprobs": USE_ROLLOUT_LOGPS,
     "max_tokens_to_generate": MAX_RESPONSE_LENGTH,
 }
@@ -911,9 +911,9 @@ def get_lora_model(base_model, model_mesh):
   }
   # Force no-op mappings for weight sync if both trainer and sampler use MaxText
   if hasattr(qwen_reference, "use_no_op_mappings"):
-    qwen_reference.use_no_op_mappings = True
+    qwen_reference.use_no_op_mappings = True  # pyrefly: ignore[missing-attribute]
   if hasattr(qwen_actor, "use_no_op_mappings"):
-    qwen_actor.use_no_op_mappings = True
+    qwen_actor.use_no_op_mappings = True  # pyrefly: ignore[missing-attribute]
     logging.info("Forced use_no_op_mappings=True on actor/reference models.")
 
 
diff --git a/examples/frozenlake/agent.py b/examples/frozenlake/agent.py
index 197617381..b4e828dce 100644
--- a/examples/frozenlake/agent.py
+++ b/examples/frozenlake/agent.py
@@ -198,7 +198,7 @@ def update_from_model(self, response: str, **kwargs) -> agent_types.Action:
     # Record complete step with conversation context and parsed action.
     cur_step = self._trajectory.steps[-1]
     cur_step.thought = thought
-    cur_step.action = action_str
+    cur_step.action = action_str  # pyrefly: ignore[bad-assignment]
     cur_step.model_response = response
 
     self.step += 1
diff --git a/examples/frozenlake/data.py b/examples/frozenlake/data.py
index 9e41b18c1..e80d86328 100644
--- a/examples/frozenlake/data.py
+++ b/examples/frozenlake/data.py
@@ -21,8 +21,8 @@
 
 import numpy as np
 import pandas as pd
-import datasets as datasets_lib
-import grain
+import datasets as datasets_lib  # pyrefly: ignore[missing-import]
+import grain  # pyrefly: ignore[missing-import]
 
 
 DEFAULT_DIR = os.getcwd()
diff --git a/examples/frozenlake/train_frozenlake.py b/examples/frozenlake/train_frozenlake.py
index 811be888e..e3db56868 100644
--- a/examples/frozenlake/train_frozenlake.py
+++ b/examples/frozenlake/train_frozenlake.py
@@ -334,8 +334,8 @@ def process_item(item):
     item["prompts"] = ""
     return item
 
-  train_ds = grain.MapDataset.source(train_ds).map(process_item)
-  test_ds = grain.MapDataset.source(test_ds).map(process_item)
+  train_ds = grain.MapDataset.source(train_ds).map(process_item)  # pyrefly: ignore[bad-argument-type]
+  test_ds = grain.MapDataset.source(test_ds).map(process_item)  # pyrefly: ignore[bad-argument-type]
   return train_ds, test_ds
 
 
@@ -349,7 +349,7 @@ def process_item(item):
 train_dataset, test_dataset = create_datasets()
 train_dataset, val_dataset = data_lib.post_init_dataset(
     train_dataset,
-    tokenizer,
+    tokenizer,  # pyrefly: ignore[bad-argument-type]
     batch_size=BATCH_SIZE,
     num_batches=NUM_BATCHES,
     max_prompt_length=MAX_PROMPT_LENGTH,
@@ -358,7 +358,7 @@ def process_item(item):
 )
 test_dataset, _ = data_lib.post_init_dataset(
     test_dataset,
-    tokenizer,
+    tokenizer,  # pyrefly: ignore[bad-argument-type]
     batch_size=BATCH_SIZE,
     num_batches=NUM_TEST_BATCHES,
     max_prompt_length=MAX_PROMPT_LENGTH,
diff --git a/examples/math_gsm8k/gemma_grpo_demo_nb.py b/examples/math_gsm8k/gemma_grpo_demo_nb.py
index 2e016f54c..9091d1697 100644
--- a/examples/math_gsm8k/gemma_grpo_demo_nb.py
+++ b/examples/math_gsm8k/gemma_grpo_demo_nb.py
@@ -291,7 +291,7 @@ def extract_hash_answer(text: str) -> str | None:
 def get_ref_model():
   """Loads the reference model, from CNS in g3 or Kaggle in OSS."""
   mesh = jax.make_mesh(
-      *MESH, axis_types=(jax.sharding.AxisType.Auto,) * len(MESH[0])
+      *MESH, axis_types=(jax.sharding.AxisType.Auto,) * len(MESH[0])  # pyrefly: ignore[bad-argument-type]
   )
 
   if ENV == 'g3':
diff --git a/examples/math_gsm8k/qwen3_grpo_demo.py b/examples/math_gsm8k/qwen3_grpo_demo.py
index 83c04a305..9734cf12d 100644
--- a/examples/math_gsm8k/qwen3_grpo_demo.py
+++ b/examples/math_gsm8k/qwen3_grpo_demo.py
@@ -625,7 +625,7 @@ def main() -> None:
       trust_remote_code=True,
   )
   chat_parser = VTCRawTextParser()
-  qwen_eos_tokens = tokenizer.encode("<|im_end|>", add_special_tokens=False)
+  qwen_eos_tokens = tokenizer.encode("<|im_end|>", add_special_tokens=False)  # pyrefly: ignore[missing-attribute]
 
   reference, actor = create_reference_and_actor(shared_mesh)
   show_hbm_usage("after loading qwen_ref / qwen_actor")
diff --git a/tests/experimental/train/peft_trainer_v2_test.py b/tests/experimental/train/peft_trainer_v2_test.py
index 51e38a407..5501ee8af 100644
--- a/tests/experimental/train/peft_trainer_v2_test.py
+++ b/tests/experimental/train/peft_trainer_v2_test.py
@@ -149,19 +149,19 @@ def test_basic_training(self, cache_nnx_graph: bool):
     jax.tree.map_with_path(tc.assert_not_equal, original_variables, variables)
 
     self.assertGreater(
-        trainer.metrics_logger.get_metric('', 'perplexity', 'train'), 0
+        trainer.metrics_logger.get_metric('', 'perplexity', 'train'), 0  # pyrefly: ignore[missing-attribute]
     )
     self.assertEqual(
-        trainer.metrics_logger.get_metric('', 'learning_rate', 'train'),
+        trainer.metrics_logger.get_metric('', 'learning_rate', 'train'),  # pyrefly: ignore[missing-attribute]
         TEST_LEARNING_RATE,
     )
     self.assertGreater(
-        trainer.metrics_logger.get_metric('', 'perplexity', 'eval'), 0
+        trainer.metrics_logger.get_metric('', 'perplexity', 'eval'), 0  # pyrefly: ignore[missing-attribute]
     )
     self.assertGreater(trainer._train_steps, 0)
 
     self.assertLen(
-        trainer.metrics_logger.get_metric_history('', 'perplexity', 'train'),
+        trainer.metrics_logger.get_metric_history('', 'perplexity', 'train'),  # pyrefly: ignore[missing-attribute]
         trainer._train_steps,
     )
 
@@ -416,7 +416,7 @@ def test_lora_training(self, learning_rate_scheduler):
         tc.assert_not_equal, original_lora_params, lora_params
     )
     self.assertEqual(
-        trainer.metrics_logger.get_metric('', 'learning_rate', 'train'),
+        trainer.metrics_logger.get_metric('', 'learning_rate', 'train'),  # pyrefly: ignore[missing-attribute]
         TEST_LEARNING_RATE,
     )
 
@@ -446,7 +446,7 @@ def train(
 
       trainer.train(train_ds, self.eval_ds)
       self.assertEqual(
-          trainer.metrics_logger.get_metric('', 'learning_rate', 'train'),
+          trainer.metrics_logger.get_metric('', 'learning_rate', 'train'),  # pyrefly: ignore[missing-attribute]
           TEST_LEARNING_RATE,
       )
       return nnx.state(model, nnx.Param), trainer
@@ -675,7 +675,7 @@ def _post_process_train_step(self, aux):
         if self._buffered_train_metrics is not None:
           self._buffered_train_metrics.additional_metrics['foo'] = (
               [aux['foo']],
-              lambda xs: xs[-1],
+              lambda xs: xs[-1],  # pyrefly: ignore[bad-index]
           )
 
       def _post_process_eval_step(self, aux):
@@ -686,7 +686,7 @@ def _post_process_eval_step(self, aux):
         if self._buffered_eval_metrics is not None:
           self._buffered_eval_metrics.additional_metrics['foo'] = (
               [aux['foo']],
-              lambda xs: xs[-1],
+              lambda xs: xs[-1],  # pyrefly: ignore[bad-index]
           )
 
     config = peft_trainer_v2.TrainingConfig(eval_every_n_steps=2, max_steps=100)
@@ -732,7 +732,7 @@ def test_get_metrics(self):
     self.assertEqual(metrics.id, 0)
     self.assertEqual(metrics.mode, 'eval')
     self.assertIn('loss', metrics.scalar_metrics)
-    self.assertGreater(metrics.scalar_metrics['loss'], 0)
+    self.assertGreater(metrics.scalar_metrics['loss'], 0)  # pyrefly: ignore[no-matching-overload]
 
     # After calling get_metrics, the buffer should be cleared
     self.assertEqual(trainer.get_metrics().id, -1)
@@ -744,7 +744,7 @@ def test_get_metrics(self):
     self.assertEqual(train_metrics.mode, 'train')
     self.assertIn('loss', train_metrics.scalar_metrics)
     self.assertIn('grad_norm', train_metrics.scalar_metrics)
-    self.assertGreater(train_metrics.scalar_metrics['loss'], 0)
+    self.assertGreater(train_metrics.scalar_metrics['loss'], 0)  # pyrefly: ignore[no-matching-overload]
 
   def test_injected_params(self):
     config = peft_trainer_v2.TrainingConfig(eval_every_n_steps=2, max_steps=100)
@@ -760,7 +760,7 @@ def test_injected_params(self):
     trainer = trainer.with_gen_model_input_fn(dummy_gen_model_input_fn)
     trainer.train(self.train_ds, self.eval_ds)
     self.assertEqual(
-        trainer.metrics_logger.get_metric('', 'learning_rate', 'train'),
+        trainer.metrics_logger.get_metric('', 'learning_rate', 'train'),  # pyrefly: ignore[missing-attribute]
         TEST_LEARNING_RATE,
     )
 
diff --git a/tests/generate/sampler_test.py b/tests/generate/sampler_test.py
index 2b3d9fa73..f1ed6b4ba 100644
--- a/tests/generate/sampler_test.py
+++ b/tests/generate/sampler_test.py
@@ -431,7 +431,7 @@ def test_decode_stops_after_prefill_for_single_generation_step(self):
     prompt_tokens = sampler.tokenize('input string')
     all_input_ids = jnp.array([
         utils.pad_to_length(
-            prompt_tokens,
+            prompt_tokens,  # pyrefly: ignore[bad-argument-type]
             target_length=max_prompt_length,
             pad_value=vocab.pad_id(),
             left=True,
diff --git a/tests/sft/otel_wandb_test.py b/tests/sft/otel_wandb_test.py
index f30339e64..2f7bf4847 100644
--- a/tests/sft/otel_wandb_test.py
+++ b/tests/sft/otel_wandb_test.py
@@ -35,7 +35,7 @@ class WandbMetricsExporterTest(absltest.TestCase):
 
   def setUp(self):
     super().setUp()
-    self.run = _FakeWandbRun()
+    self.run = _FakeWandbRun()  # pyrefly: ignore[bad-assignment]
     self.exporter = otel_wandb.WandbMetricsExporter(self.run)
     self.reader = otel_sdk_export.PeriodicExportingMetricReader(
         self.exporter, export_interval_millis=3_600_000
@@ -76,8 +76,8 @@ def test_double_write_reaches_wandb_run(self):
 
     self.meter_provider.force_flush()
 
-    self.assertLen(self.run.calls, 1)
-    values, step = self.run.calls[0]
+    self.assertLen(self.run.calls, 1)  # pyrefly: ignore[missing-attribute]
+    values, step = self.run.calls[0]  # pyrefly: ignore[missing-attribute]
     self.assertEqual(step, 3)
     self.assertAlmostEqual(values["actor/train/tunix.training.loss"], 0.5)
     self.assertAlmostEqual(
@@ -91,10 +91,10 @@ def test_groups_are_logged_in_step_order(self):
 
     self.meter_provider.force_flush()
 
-    steps = [step for _, step in self.run.calls]
+    steps = [step for _, step in self.run.calls]  # pyrefly: ignore[missing-attribute]
     self.assertEqual(steps, sorted(steps))
     logged_keys = set()
-    for values, _ in self.run.calls:
+    for values, _ in self.run.calls:  # pyrefly: ignore[missing-attribute]
       logged_keys.update(values)
     self.assertIn("actor/train/tunix.training.loss", logged_keys)
     self.assertIn("critic/eval/tunix.training.loss", logged_keys)
diff --git a/tests/utils/mesh_test.py b/tests/utils/mesh_test.py
index f1c837943..4a5a59604 100644
--- a/tests/utils/mesh_test.py
+++ b/tests/utils/mesh_test.py
@@ -499,7 +499,7 @@ def __init__(self, device_id, coords):
 
     allocated, _ = mesh._allocate_devices_by_coords(fake_devices, 2)
 
-    self.assertEqual([device.id for device in allocated], [0, 1])
+    self.assertEqual([device.id for device in allocated], [0, 1])  # pyrefly: ignore[not-iterable]
 
   def test_allocate_named_mesh_device_slices_prefers_coord_boxes(self):
     class FakeDevice:
@@ -701,7 +701,7 @@ def __init__(self, device_id, coords):
         allocation_policy="COMPACT",
     )
 
-    allocated_coords = [device.coords for device in allocated]
+    allocated_coords = [device.coords for device in allocated]  # pyrefly: ignore[not-iterable]
     mins = tuple(
         min(coords[dim] for coords in allocated_coords) for dim in range(3)
     )
diff --git a/tunix/cli/base_rl_pipeline.py b/tunix/cli/base_rl_pipeline.py
index cf21a6c42..c8c4fdb47 100644
--- a/tunix/cli/base_rl_pipeline.py
+++ b/tunix/cli/base_rl_pipeline.py
@@ -589,7 +589,7 @@ def create_rl_engine(self, tokenizer) -> rl_engine_lib.RLEngine:
           params=jax.random.key(critic_model_config.get("rng_seed", 0))
       )
 
-      if hasattr(critic_model.config.shd_config, "score_weight_d1"):
+      if hasattr(critic_model.config.shd_config, "score_weight_d1"):  # pyrefly: ignore[missing-attribute]
         critic_model = rl_utils.TransformerWithScoreHead(
             critic_model, rngs=rngs
         )
@@ -886,7 +886,7 @@ def _run(self, mode: str):
   def run_trainer(self):
     """Dispatch to standard or agentic trainer based on training_mode."""
     mode = self.config.get("training_mode", self._default_training_mode)
-    self._run(mode=mode)
+    self._run(mode=mode)  # pyrefly: ignore[bad-argument-type]
 
 
 def setup_jax_pathways(pathways_bns: str):
diff --git a/tunix/cli/config.py b/tunix/cli/config.py
index 8a2188a6c..662aefaf6 100644
--- a/tunix/cli/config.py
+++ b/tunix/cli/config.py
@@ -283,7 +283,7 @@ def _config_mapping(self, key: str) -> dict[str, Any]:
           f" {type(value).__name__}."
       )
     if isinstance(value, omegaconf.DictConfig):
-      return omegaconf.OmegaConf.to_container(value, resolve=True)
+      return omegaconf.OmegaConf.to_container(value, resolve=True)  # pyrefly: ignore[bad-return]
     return dict(value)
 
   def _mutable_config_mapping(self, key: str) -> MutableMapping[str, Any]:
diff --git a/tunix/examples/data/translation_dataset.py b/tunix/examples/data/translation_dataset.py
index 2e1bf1bbb..ccdbae72f 100644
--- a/tunix/examples/data/translation_dataset.py
+++ b/tunix/examples/data/translation_dataset.py
@@ -69,7 +69,7 @@ def create_datasets(
         dataset_name, split=("train", "valid"), download=tfds_download
     )
   elif dataset_name == "Helsinki-NLP/opus-100":  # Hugging Face dataloader
-    train_ds, eval_ds = datasets.load_dataset(
+    train_ds, eval_ds = datasets.load_dataset(  # pyrefly: ignore[no-matching-overload]
         dataset_name, data_dir="en-fr", split=("train", "validation")  # pyrefly: ignore[bad-argument-type]
     )
   else:
diff --git a/tunix/experimental/distributed/examples/basics/door.py b/tunix/experimental/distributed/examples/basics/door.py
index d071f001f..942059234 100644
--- a/tunix/experimental/distributed/examples/basics/door.py
+++ b/tunix/experimental/distributed/examples/basics/door.py
@@ -17,7 +17,7 @@ def main(argv, context: ProcessContext | None) -> None:
   knock_future = futures.Future()
   assert context is not None
   context.ipc.discovery.on_register(
-      callback=lambda hostname, _, metadata: (
+      callback=lambda hostname, _, metadata: (  # pyrefly: ignore[bad-argument-type]
           logging.info(
               f"{hostname} knocked and said: {pickle.loads(metadata)}"
           ),
diff --git a/tunix/experimental/distributed/examples/vllm_rollout/rollout.py b/tunix/experimental/distributed/examples/vllm_rollout/rollout.py
index 8a604e188..f5f7a6b5e 100644
--- a/tunix/experimental/distributed/examples/vllm_rollout/rollout.py
+++ b/tunix/experimental/distributed/examples/vllm_rollout/rollout.py
@@ -64,7 +64,7 @@ def main(argv: Sequence[str], context: ProcessContext | None) -> None:
 
   tokenizer = AutoTokenizer.from_pretrained(args.model_name)
   config = vllm_sampler.VllmConfig(engine_kwargs={"model": args.model_name})
-  sampler_server = legacy_sampler_lib.LegacyVllmSamplerAdapter(
+  sampler_server = legacy_sampler_lib.LegacyVllmSamplerAdapter(  # pyrefly: ignore[bad-instantiation]
       server_id="vllm-0",
       tokenizer=tokenizer,
       config=config,
diff --git a/tunix/experimental/orchestrator/async_rl_program.py b/tunix/experimental/orchestrator/async_rl_program.py
index 6210529bf..2364eb7f6 100644
--- a/tunix/experimental/orchestrator/async_rl_program.py
+++ b/tunix/experimental/orchestrator/async_rl_program.py
@@ -165,9 +165,9 @@ async def critique_stage(
       for idx, payload in enumerate(trainer_payloads):
         adv = payload.advantages
         reward_val = (
-            float(adv[0])
-            if hasattr(adv, "__len__") and len(adv) > 0
-            else float(adv)
+            float(adv[0])  # pyrefly: ignore[bad-index]
+            if hasattr(adv, "__len__") and len(adv) > 0  # pyrefly: ignore[bad-argument-type]
+            else float(adv)  # pyrefly: ignore[bad-argument-type]
         )
         item = datatypes.TrajectoryItem(
             pair_index=idx,
@@ -176,7 +176,7 @@ async def critique_stage(
             traj=datatypes.Trajectory(reward=reward_val),
             # TODO: Stream RLTrainerPayload directly instead of re-wrapping in TrajectoryItem.
         )
-        item.payload = payload
+        item.payload = payload  # pyrefly: ignore[missing-attribute]
         await self.scored_q.put(item)
 
   async def train_stage(
@@ -199,7 +199,7 @@ async def train_stage(
 
         payloads = [getattr(item, "payload", None) for item in scored_items]
         # TODO: Implement streaming microbatch assembly to overlap packing with trainer execution.
-        microbatches = self.assembler.pack(payloads)
+        microbatches = self.assembler.pack(payloads)  # pyrefly: ignore[bad-argument-type]
 
         is_final = group_idx == self.mini_batch_size - 1
         for batch in microbatches:
@@ -243,9 +243,9 @@ async def run_async(
         )
         for task in done:
           if task.exception():
-            raise task.exception()
+            raise task.exception()  # pyrefly: ignore[bad-raise]
       if train_task.exception():
-        raise train_task.exception()
+        raise train_task.exception()  # pyrefly: ignore[bad-raise]
     except Exception as exc:
       logging.error("Exception in StandardRLProgram execution: %s", exc)
       await self.raw_q.abort(exc)
diff --git a/tunix/experimental/orchestrator/batch_assembly.py b/tunix/experimental/orchestrator/batch_assembly.py
index 08baaf981..7173a1da6 100644
--- a/tunix/experimental/orchestrator/batch_assembly.py
+++ b/tunix/experimental/orchestrator/batch_assembly.py
@@ -52,7 +52,7 @@ def pack(self, items: Sequence[datatypes.RLTrainerPayload]) -> list[datatypes.RL
     # Calculate token lengths from explicit fields
     item_lengths = []
     for it in items:
-      item_lengths.append(len(it.token_ids) if it.token_ids is not None else 0)
+      item_lengths.append(len(it.token_ids) if it.token_ids is not None else 0)  # pyrefly: ignore[bad-argument-type]
 
     item_list = sorted(zip(items, item_lengths), key=lambda x: x[1], reverse=True)
 
diff --git a/tunix/experimental/orchestrator/distributed_rl_engine.py b/tunix/experimental/orchestrator/distributed_rl_engine.py
index 7a884290a..91736b745 100644
--- a/tunix/experimental/orchestrator/distributed_rl_engine.py
+++ b/tunix/experimental/orchestrator/distributed_rl_engine.py
@@ -290,7 +290,7 @@ async def train_step(
         **kwargs,
     )
 
-  async def sync_weights(
+  async def sync_weights(  # pyrefly: ignore[bad-override]
       self,
       role: datatypes.Role = datatypes.Role.ACTOR,
       target_roles: Sequence[datatypes.Role] | None = None,
diff --git a/tunix/experimental/orchestrator/health_monitor.py b/tunix/experimental/orchestrator/health_monitor.py
index d6e59eef2..3b3db9cfd 100644
--- a/tunix/experimental/orchestrator/health_monitor.py
+++ b/tunix/experimental/orchestrator/health_monitor.py
@@ -151,7 +151,7 @@ def overdue(self) -> list[OverdueWorker]:
     now = self._clock()
     result: list[OverdueWorker] = []
     for worker_id, (state, since) in sorted(self._state_since.items()):
-      deadline = self._deadlines.get(state)
+      deadline = self._deadlines.get(state)  # pyrefly: ignore[bad-argument-type]
       if deadline is None:
         continue
       elapsed = now - since
diff --git a/tunix/experimental/orchestrator/orchestrator.py b/tunix/experimental/orchestrator/orchestrator.py
index a32587f39..f97b17506 100644
--- a/tunix/experimental/orchestrator/orchestrator.py
+++ b/tunix/experimental/orchestrator/orchestrator.py
@@ -165,7 +165,7 @@ def run(
         assembler=active_assembler,
     )
     self.run_program(
-        program=active_program,
+        program=active_program,  # pyrefly: ignore[bad-argument-type]
         train_dataset=dataset,
         num_steps=num_steps,
         bring_up=False,
diff --git a/tunix/experimental/queue_manager/trajectory_queue_manager.py b/tunix/experimental/queue_manager/trajectory_queue_manager.py
index fc3ede461..d8440ee2d 100644
--- a/tunix/experimental/queue_manager/trajectory_queue_manager.py
+++ b/tunix/experimental/queue_manager/trajectory_queue_manager.py
@@ -110,7 +110,7 @@ async def get_batch(
         out.extend(g)
       return out
     actual_batch_size = batch_size if batch_size is not None else self.group_size
-    return await super().get_batch(batch_size=actual_batch_size)
+    return await super().get_batch(batch_size=actual_batch_size)  # pyrefly: ignore[bad-argument-type]
 
   def commit(self, step: int, groups: Sequence[Any] | None = None) -> None:
     """Commits in-flight groups after a successful global step boundary."""
diff --git a/tunix/experimental/rollout/collector.py b/tunix/experimental/rollout/collector.py
index b7837425a..db6058b3c 100644
--- a/tunix/experimental/rollout/collector.py
+++ b/tunix/experimental/rollout/collector.py
@@ -105,7 +105,7 @@ async def model_call(
     inner_engine = rl_collect_engine.TrajectoryCollectEngine(
         agent=self.agent,
         env=self.env,
-        model_call=model_call,
+        model_call=model_call,  # pyrefly: ignore[bad-argument-type]
         tokenizer=self.tokenizer,
         chat_parser=self.chat_parser,
     )
diff --git a/tunix/experimental/rollout/manager.py b/tunix/experimental/rollout/manager.py
index d369caa45..e2478e193 100644
--- a/tunix/experimental/rollout/manager.py
+++ b/tunix/experimental/rollout/manager.py
@@ -57,11 +57,11 @@ def __init__(
       elif sampler_type == "legacy_vllm":
         from tunix.experimental.rollout import legacy_vllm_sampler_adapter  # pylint: disable=g-import-not-at-top
 
-        sampler = legacy_vllm_sampler_adapter.LegacyVllmSamplerAdapter(
+        sampler = legacy_vllm_sampler_adapter.LegacyVllmSamplerAdapter(  # pyrefly: ignore[bad-instantiation]
             server_id="legacy_vllm_sampler",
         )
       elif sampler_type == "vanilla":
-        sampler = vanilla_sampler_adapter.VanillaSamplerAdapter(
+        sampler = vanilla_sampler_adapter.VanillaSamplerAdapter(  # pyrefly: ignore[bad-instantiation]
             server_id="vanilla_sampler",
         )
       else:
@@ -164,7 +164,7 @@ async def generate(
           for req in requests
       ]
       return await asyncio.gather(*tasks)
-    return await self._generate_one(requests, on_complete=on_complete)
+    return await self._generate_one(requests, on_complete=on_complete)  # pyrefly: ignore[bad-argument-type]
 
   async def _run_and_enqueue(
       self,
diff --git a/tunix/experimental/train/peft_trainer_v2.py b/tunix/experimental/train/peft_trainer_v2.py
index e2202fe18..df994c284 100644
--- a/tunix/experimental/train/peft_trainer_v2.py
+++ b/tunix/experimental/train/peft_trainer_v2.py
@@ -562,7 +562,7 @@ def diff_fn(model, *args, **kwargs):
       if isinstance(out, utils.LossOutput):
         return out.primary_loss.unreduced_sum, out
       elif self._has_aux:
-        return out[0], out[1]
+        return out[0], out[1]  # pyrefly: ignore[bad-index]
       else:
         return out, None
 
@@ -693,7 +693,7 @@ def _shard(x, p):
         return x
       if p is None:
         p = shd.PartitionSpec()
-      sharding = sharding_utils.get_sharding(x, mesh, p)
+      sharding = sharding_utils.get_sharding(x, mesh, p)  # pyrefly: ignore[bad-argument-type]
       if hasattr(x, "sharding") and x.sharding == sharding:
         return x
       if getattr(x, "is_fully_addressable", True):
@@ -1278,7 +1278,7 @@ def train(
           span.device_end([computation_to_track])
           span_v2.async_end([computation_to_track])
 
-        self._throttler.add_computation(computation_to_track)
+        self._throttler.add_computation(computation_to_track)  # pyrefly: ignore[bad-argument-type]
         if is_update_step_val:
           self.save_checkpoint()
           if (
diff --git a/tunix/experimental/worker/remote_execution.py b/tunix/experimental/worker/remote_execution.py
index 4c3f22635..60b0c4207 100644
--- a/tunix/experimental/worker/remote_execution.py
+++ b/tunix/experimental/worker/remote_execution.py
@@ -905,7 +905,7 @@ def _get_next_actor(
       ):
         return getattr(self.router, method_name)(self._actors, args, kwargs)
       elif callable(self.router):
-        return self.router(self._actors, method_name, args, kwargs)
+        return self.router(self._actors, method_name, args, kwargs)  # pyrefly: ignore[bad-return]
       else:
         raise TypeError(
             f"Router object {type(self.router)} must provide a method matching "
@@ -1259,7 +1259,7 @@ def execute(self, *f_args, **f_kwargs):
             pass
         return handle.submit("execute", *args, **kwargs)
 
-      remote_func.remote = remote_func
+      remote_func.remote = remote_func  # pyrefly: ignore[missing-attribute]
       return remote_func
     else:
       raise TypeError(
diff --git a/tunix/experimental/worker/rollout_worker.py b/tunix/experimental/worker/rollout_worker.py
index 2b0fc90a2..4386e4961 100644
--- a/tunix/experimental/worker/rollout_worker.py
+++ b/tunix/experimental/worker/rollout_worker.py
@@ -158,7 +158,7 @@ def _to_rollout_response(
           or getattr(item, "trajectory_id", "")
           or getattr(item, "prompt_id", ""),
           status="ERROR",
-          error=item.error_message,
+          error=item.error_message,  # pyrefly: ignore[bad-argument-type]
           prompt_tokens=(
               prompt_tokens
               if prompt_tokens is not None
@@ -170,7 +170,7 @@ def _to_rollout_response(
       req_id = request_id or getattr(item, "trajectory_id", "default")
       return datatypes.RolloutResponse.from_trajectory(
           request_id=req_id,
-          traj=item,
+          traj=item,  # pyrefly: ignore[bad-argument-type]
           prompt_tokens=(
               prompt_tokens
               if prompt_tokens is not None
diff --git a/tunix/generate/vllm_sampler.py b/tunix/generate/vllm_sampler.py
index efc750b87..634853c64 100644
--- a/tunix/generate/vllm_sampler.py
+++ b/tunix/generate/vllm_sampler.py
@@ -286,7 +286,7 @@ def _vllm_config(self, config: VllmConfig):
       args["additional_config"]["lora_config"] = config.lora_config
 
     tp, dp, ep = utils.resolve_parallelism_sizes(
-        mesh=config.mesh,
+        mesh=config.mesh,  # pyrefly: ignore[bad-argument-type]
         tensor_parallel_size=config.tensor_parallel_size,
         data_parallel_size=config.data_parallel_size,
         expert_parallel_size=config.expert_parallel_size,
@@ -377,10 +377,10 @@ def detokenize(
             np.array(single_output.token_ids, dtype=np.int32)
         )
         decoded_outputs[idx].append(
-            self.tokenizer.decode(single_output.token_ids)
+            self.tokenizer.decode(single_output.token_ids)  # pyrefly: ignore[bad-argument-type]
         )
         logprobs = utils.get_logprobs_from_vllm_output(
-            list(single_output.token_ids), single_output.logprobs
+            list(single_output.token_ids), single_output.logprobs  # pyrefly: ignore[bad-argument-type]
         )
         out_logprobs[idx].append(logprobs)
         logging.debug(
@@ -470,7 +470,7 @@ def __call__(
         else:
           sampling_params = SamplingParams()
       else:
-        sampling_params = self.llm.get_default_sampling_params()
+        sampling_params = self.llm.get_default_sampling_params()  # pyrefly: ignore[missing-attribute]
       sampling_params.detokenize = False
       sampling_params.max_tokens = max_generation_steps
       sampling_params.n = multi_sampling
@@ -532,7 +532,7 @@ def __call__(
     if self._driver is not None:
       outputs = self._generate_server_mode(prompt_objects, sampling_params)
     else:
-      outputs = self.llm.generate(
+      outputs = self.llm.generate(  # pyrefly: ignore[missing-attribute]
           prompts=prompt_objects,
           sampling_params=sampling_params,
           use_tqdm=True,
@@ -566,5 +566,5 @@ def __call__(
         logits=None,
         tokens=out_tokens[0],
         padded_prompt_tokens=all_input_ids,
-        logprobs=out_logprobs[0] if self.config.return_logprobs else None,
+        logprobs=out_logprobs[0] if self.config.return_logprobs else None,  # pyrefly: ignore[bad-argument-type]
     )
diff --git a/tunix/models/automodel.py b/tunix/models/automodel.py
index 5c41a6318..5e47d6ffd 100644
--- a/tunix/models/automodel.py
+++ b/tunix/models/automodel.py
@@ -571,7 +571,7 @@ def from_pretrained(
       # Get load_dtype explicitly from kwargs
       load_dtype_str = kwargs.get('load_dtype')
       try:
-        load_dtype = getattr(jnp, load_dtype_str)
+        load_dtype = getattr(jnp, load_dtype_str)  # pyrefly: ignore[bad-argument-type]
       except AttributeError:
         raise ValueError(
             f"Invalid load_dtype: {load_dtype_str}. Must be a valid"
@@ -586,7 +586,7 @@ def from_pretrained(
         valid_fields = {f.name for f in dataclasses.fields(model_params)}
         overrides = {k: v for k, v in kwargs.items() if k in valid_fields and v is not None}
         if 'remat_config' in overrides and isinstance(overrides['remat_config'], str):
-          model_module = get_model_module(naming_info.model_name, ModelModule.MODEL)
+          model_module = get_model_module(naming_info.model_name, ModelModule.MODEL)  # pyrefly: ignore[bad-argument-type]
           if hasattr(model_module, 'RematConfig'):
             remat_cfg_str = overrides['remat_config']
             try:
diff --git a/tunix/models/gemma/model.py b/tunix/models/gemma/model.py
index 0b163ee69..10414bda5 100644
--- a/tunix/models/gemma/model.py
+++ b/tunix/models/gemma/model.py
@@ -73,18 +73,18 @@ def get_default_sharding(is_sampling: bool = False):
     fsdp = 'fsdp' if not is_sampling else None
 
     return ShardingConfig(
-        emb_vd=P('tp', fsdp),
-        q_weight_ndh=P('tp', fsdp, None),
-        kv_weight_cndh=P(None, 'tp', fsdp, None),
-        qkv_weight_cndh=P(None, 'tp', fsdp, None),
-        o_weight_nhd=P('tp', None, fsdp),
-        ffw_weight_df=P(fsdp, 'tp'),
-        ffw_weight_fd=P('tp', fsdp),
-        rms_norm_weight=P('tp',),
-        act_btd=P('fsdp', None, None if is_sampling else 'tp'),
-        act_btf=P('fsdp', None, 'tp'),
-        act_btnh=P('fsdp', None, 'tp', None),
-        score_weight_d1=P(fsdp, None),
+        emb_vd=P('tp', fsdp),  # pyrefly: ignore[bad-argument-type]
+        q_weight_ndh=P('tp', fsdp, None),  # pyrefly: ignore[bad-argument-type]
+        kv_weight_cndh=P(None, 'tp', fsdp, None),  # pyrefly: ignore[bad-argument-type]
+        qkv_weight_cndh=P(None, 'tp', fsdp, None),  # pyrefly: ignore[bad-argument-type]
+        o_weight_nhd=P('tp', None, fsdp),  # pyrefly: ignore[bad-argument-type]
+        ffw_weight_df=P(fsdp, 'tp'),  # pyrefly: ignore[bad-argument-type]
+        ffw_weight_fd=P('tp', fsdp),  # pyrefly: ignore[bad-argument-type]
+        rms_norm_weight=P('tp',),  # pyrefly: ignore[bad-argument-type]
+        act_btd=P('fsdp', None, None if is_sampling else 'tp'),  # pyrefly: ignore[bad-argument-type]
+        act_btf=P('fsdp', None, 'tp'),  # pyrefly: ignore[bad-argument-type]
+        act_btnh=P('fsdp', None, 'tp', None),  # pyrefly: ignore[bad-argument-type]
+        score_weight_d1=P(fsdp, None),  # pyrefly: ignore[bad-argument-type]
     )
 
 
diff --git a/tunix/models/gemma3/model.py b/tunix/models/gemma3/model.py
index d660157d9..c5b7d9600 100644
--- a/tunix/models/gemma3/model.py
+++ b/tunix/models/gemma3/model.py
@@ -72,20 +72,20 @@ def get_default_sharding(is_sampling: bool = False):
     fsdp = 'fsdp' if not is_sampling else None
 
     return ShardingConfig(
-        emb_vd=P('tp', fsdp),
-        q_weight_ndh=P('tp', fsdp, None),
-        kv_weight_cndh=P(None, 'tp', fsdp, None),
-        qkv_weight_cndh=P(None, 'tp', fsdp, None),
-        o_weight_nhd=P('tp', None, fsdp),
-        ffw_weight_df=P(fsdp, 'tp'),
-        ffw_weight_fd=P('tp', fsdp),
-        rms_norm_weight=P('tp',),
-        act_btd=P('fsdp', None, None if is_sampling else 'tp'),
-        act_btf=P('fsdp', None, 'tp'),
-        act_btnh=P('fsdp', None, 'tp', None),
-        score_weight_d1=P(fsdp, None),
-        vision_proj=P(fsdp, 'tp'),
-        vision_soft_emb_norm_weight=P('tp',),
+        emb_vd=P('tp', fsdp),  # pyrefly: ignore[bad-argument-type]
+        q_weight_ndh=P('tp', fsdp, None),  # pyrefly: ignore[bad-argument-type]
+        kv_weight_cndh=P(None, 'tp', fsdp, None),  # pyrefly: ignore[bad-argument-type]
+        qkv_weight_cndh=P(None, 'tp', fsdp, None),  # pyrefly: ignore[bad-argument-type]
+        o_weight_nhd=P('tp', None, fsdp),  # pyrefly: ignore[bad-argument-type]
+        ffw_weight_df=P(fsdp, 'tp'),  # pyrefly: ignore[bad-argument-type]
+        ffw_weight_fd=P('tp', fsdp),  # pyrefly: ignore[bad-argument-type]
+        rms_norm_weight=P('tp',),  # pyrefly: ignore[bad-argument-type]
+        act_btd=P('fsdp', None, None if is_sampling else 'tp'),  # pyrefly: ignore[bad-argument-type]
+        act_btf=P('fsdp', None, 'tp'),  # pyrefly: ignore[bad-argument-type]
+        act_btnh=P('fsdp', None, 'tp', None),  # pyrefly: ignore[bad-argument-type]
+        score_weight_d1=P(fsdp, None),  # pyrefly: ignore[bad-argument-type]
+        vision_proj=P(fsdp, 'tp'),  # pyrefly: ignore[bad-argument-type]
+        vision_soft_emb_norm_weight=P('tp',),  # pyrefly: ignore[bad-argument-type]
         siglip=vision.SigLIPShardingConfig.get_default_sharding(is_sampling),
     )
 
diff --git a/tunix/models/gemma4/mapping_vllm_jax.py b/tunix/models/gemma4/mapping_vllm_jax.py
index b80a905ac..81f3976fb 100644
--- a/tunix/models/gemma4/mapping_vllm_jax.py
+++ b/tunix/models/gemma4/mapping_vllm_jax.py
@@ -215,8 +215,8 @@ def preprocess_src_state(src_state: Any) -> Any:
         new_flat_state.append((k_keys, k_param))
       else:
         # KV-shared layer
-        k_val = jnp.zeros_like(sample_kv_val[0])
-        v_val = jnp.zeros_like(sample_kv_val[1])
+        k_val = jnp.zeros_like(sample_kv_val[0])  # pyrefly: ignore[unsupported-operation]
+        v_val = jnp.zeros_like(sample_kv_val[1])  # pyrefly: ignore[unsupported-operation]
         k_val_t = jnp.reshape(
             jnp.transpose(k_val, (1, 0, 2)), (hidden_size, -1)
         )
diff --git a/tunix/models/gemma4/model.py b/tunix/models/gemma4/model.py
index 6768f93d8..69ca6d083 100644
--- a/tunix/models/gemma4/model.py
+++ b/tunix/models/gemma4/model.py
@@ -120,27 +120,27 @@ def get_default_sharding(is_sampling: bool = False):
     fsdp = 'fsdp' if not is_sampling else None
 
     return ShardingConfig(
-        emb_vd=P('tp', fsdp),
-        q_weight_ndh=P('tp', fsdp, None),
-        kv_weight_cndh=P(None, 'tp', fsdp, None),
-        qkv_weight_cndh=P(None, 'tp', fsdp, None),
-        o_weight_nhd=P('tp', None, fsdp),
-        ffw_weight_df=P(fsdp, 'tp'),
-        ffw_weight_fd=P('tp', fsdp),
-        rms_norm_weight=P('tp',),
-        act_btd=P('fsdp', None, None if is_sampling else 'tp'),
-        act_btf=P('fsdp', None, 'tp'),
-        act_btnh=P('fsdp', None, 'tp', None),
-        score_weight_d1=P(fsdp, None),
-        vision_proj=P(fsdp, 'tp'),
-        vision_soft_emb_norm_weight=P('tp',),
-        audio_proj=P(fsdp, 'tp'),  # TODO check if good!
-        exp_weight_edf=P(fsdp, None, None, 'tp'),
-        exp_weight_efd=P(fsdp, 'tp', None),
-        per_layer_model_projection=P(fsdp, 'tp'),
-        per_layer_input_gate=P(fsdp, 'tp'),
-        per_layer_projection=P('tp', fsdp),
-        per_layer_input_embedding=P('tp', fsdp),
+        emb_vd=P('tp', fsdp),  # pyrefly: ignore[bad-argument-type]
+        q_weight_ndh=P('tp', fsdp, None),  # pyrefly: ignore[bad-argument-type]
+        kv_weight_cndh=P(None, 'tp', fsdp, None),  # pyrefly: ignore[bad-argument-type]
+        qkv_weight_cndh=P(None, 'tp', fsdp, None),  # pyrefly: ignore[bad-argument-type]
+        o_weight_nhd=P('tp', None, fsdp),  # pyrefly: ignore[bad-argument-type]
+        ffw_weight_df=P(fsdp, 'tp'),  # pyrefly: ignore[bad-argument-type]
+        ffw_weight_fd=P('tp', fsdp),  # pyrefly: ignore[bad-argument-type]
+        rms_norm_weight=P('tp',),  # pyrefly: ignore[bad-argument-type]
+        act_btd=P('fsdp', None, None if is_sampling else 'tp'),  # pyrefly: ignore[bad-argument-type]
+        act_btf=P('fsdp', None, 'tp'),  # pyrefly: ignore[bad-argument-type]
+        act_btnh=P('fsdp', None, 'tp', None),  # pyrefly: ignore[bad-argument-type]
+        score_weight_d1=P(fsdp, None),  # pyrefly: ignore[bad-argument-type]
+        vision_proj=P(fsdp, 'tp'),  # pyrefly: ignore[bad-argument-type]
+        vision_soft_emb_norm_weight=P('tp',),  # pyrefly: ignore[bad-argument-type]
+        audio_proj=P(fsdp, 'tp'),  # TODO check if good!  # pyrefly: ignore[bad-argument-type]
+        exp_weight_edf=P(fsdp, None, None, 'tp'),  # pyrefly: ignore[bad-argument-type]
+        exp_weight_efd=P(fsdp, 'tp', None),  # pyrefly: ignore[bad-argument-type]
+        per_layer_model_projection=P(fsdp, 'tp'),  # pyrefly: ignore[bad-argument-type]
+        per_layer_input_gate=P(fsdp, 'tp'),  # pyrefly: ignore[bad-argument-type]
+        per_layer_projection=P('tp', fsdp),  # pyrefly: ignore[bad-argument-type]
+        per_layer_input_embedding=P('tp', fsdp),  # pyrefly: ignore[bad-argument-type]
         vision_shd=vision.VisionShardingConfig.get_default_sharding(
             is_sampling
         ),
diff --git a/tunix/models/llama3/model.py b/tunix/models/llama3/model.py
index 140eea9f4..ab126330f 100644
--- a/tunix/models/llama3/model.py
+++ b/tunix/models/llama3/model.py
@@ -67,18 +67,18 @@ def get_default_sharding(is_sampling: bool = False):
     fsdp = 'fsdp' if not is_sampling else None
 
     return ShardingConfig(
-        emb_vd=P('tp', fsdp),
-        emb_dv=P(fsdp, 'tp'),
-        q_weight_dnh=P(fsdp, 'tp', None),
-        kv_weight_dnh=P(fsdp, 'tp', None),
-        o_weight_nhd=P('tp', None, fsdp),
-        ffw_weight_df=P(fsdp, 'tp'),
-        ffw_weight_fd=P('tp', fsdp),
-        rms_norm_weight=P('tp',),
-        act_btd=P('fsdp', None, None if is_sampling else 'tp'),
-        act_btf=P('fsdp', None, 'tp'),
-        act_btnh=P('fsdp', None, 'tp', None),
-        score_weight_d1=P(fsdp, None),
+        emb_vd=P('tp', fsdp),  # pyrefly: ignore[bad-argument-type]
+        emb_dv=P(fsdp, 'tp'),  # pyrefly: ignore[bad-argument-type]
+        q_weight_dnh=P(fsdp, 'tp', None),  # pyrefly: ignore[bad-argument-type]
+        kv_weight_dnh=P(fsdp, 'tp', None),  # pyrefly: ignore[bad-argument-type]
+        o_weight_nhd=P('tp', None, fsdp),  # pyrefly: ignore[bad-argument-type]
+        ffw_weight_df=P(fsdp, 'tp'),  # pyrefly: ignore[bad-argument-type]
+        ffw_weight_fd=P('tp', fsdp),  # pyrefly: ignore[bad-argument-type]
+        rms_norm_weight=P('tp',),  # pyrefly: ignore[bad-argument-type]
+        act_btd=P('fsdp', None, None if is_sampling else 'tp'),  # pyrefly: ignore[bad-argument-type]
+        act_btf=P('fsdp', None, 'tp'),  # pyrefly: ignore[bad-argument-type]
+        act_btnh=P('fsdp', None, 'tp', None),  # pyrefly: ignore[bad-argument-type]
+        score_weight_d1=P(fsdp, None),  # pyrefly: ignore[bad-argument-type]
     )
 
 
diff --git a/tunix/models/qwen2/model.py b/tunix/models/qwen2/model.py
index 0f162488b..b068b15f9 100644
--- a/tunix/models/qwen2/model.py
+++ b/tunix/models/qwen2/model.py
@@ -81,14 +81,14 @@ def get_default_sharding(is_sampling: bool = False, enable_sp: bool = False):
         o_weight_nhd=P('tp', None, fsdp),  # pyrefly: ignore[bad-argument-type]
         ffw_weight_df=P(fsdp, 'tp'),  # pyrefly: ignore[bad-argument-type]
         ffw_weight_fd=P('tp', fsdp),  # pyrefly: ignore[bad-argument-type]
-        rms_norm_weight=P('tp',),
-        act_btd=P('fsdp', sp, None if is_sampling else 'tp'),
-        act_btf=P('fsdp', sp, 'tp'),
-        act_btnh=P('fsdp', sp, 'tp', None),
-        score_weight_d1=P(fsdp, None),
-        exp_weight_cdf=P('fsdp', None, 'tp'),
-        exp_weight_cfd=P('fsdp', 'tp', None),
-        qkv_bias=P('tp',),
+        rms_norm_weight=P('tp',),  # pyrefly: ignore[bad-argument-type]
+        act_btd=P('fsdp', sp, None if is_sampling else 'tp'),  # pyrefly: ignore[bad-argument-type]
+        act_btf=P('fsdp', sp, 'tp'),  # pyrefly: ignore[bad-argument-type]
+        act_btnh=P('fsdp', sp, 'tp', None),  # pyrefly: ignore[bad-argument-type]
+        score_weight_d1=P(fsdp, None),  # pyrefly: ignore[bad-argument-type]
+        exp_weight_cdf=P('fsdp', None, 'tp'),  # pyrefly: ignore[bad-argument-type]
+        exp_weight_cfd=P('fsdp', 'tp', None),  # pyrefly: ignore[bad-argument-type]
+        qkv_bias=P('tp',),  # pyrefly: ignore[bad-argument-type]
     )
 
 
diff --git a/tunix/models/qwen3/model.py b/tunix/models/qwen3/model.py
index 0ebbcb3ad..aded476dc 100644
--- a/tunix/models/qwen3/model.py
+++ b/tunix/models/qwen3/model.py
@@ -115,22 +115,22 @@ def get_default_sharding(is_sampling: bool = False, enable_sp: bool = False):
     fsdp = (fsdp, sp) if fsdp and sp else fsdp
 
     return ShardingConfig(
-        emb_vd=P('tp', fsdp),
-        emb_dv=P(fsdp, 'tp'),
-        q_weight_dnh=P(fsdp, 'tp', None),
-        kv_weight_dnh=P(fsdp, 'tp', None),
-        o_weight_nhd=P('tp', None, fsdp),
-        ffw_weight_df=P(fsdp, 'tp'),
-        ffw_weight_fd=P('tp', fsdp),
-        rms_norm_weight=P(
+        emb_vd=P('tp', fsdp),  # pyrefly: ignore[bad-argument-type]
+        emb_dv=P(fsdp, 'tp'),  # pyrefly: ignore[bad-argument-type]
+        q_weight_dnh=P(fsdp, 'tp', None),  # pyrefly: ignore[bad-argument-type]
+        kv_weight_dnh=P(fsdp, 'tp', None),  # pyrefly: ignore[bad-argument-type]
+        o_weight_nhd=P('tp', None, fsdp),  # pyrefly: ignore[bad-argument-type]
+        ffw_weight_df=P(fsdp, 'tp'),  # pyrefly: ignore[bad-argument-type]
+        ffw_weight_fd=P('tp', fsdp),  # pyrefly: ignore[bad-argument-type]
+        rms_norm_weight=P(  # pyrefly: ignore[bad-argument-type]
             'tp',
         ),
-        act_btd=P('fsdp', sp, None if is_sampling else 'tp'),
-        act_btf=P('fsdp', sp, 'tp'),
-        act_btnh=P('fsdp', sp, 'tp', None),
-        score_weight_d1=P(fsdp, None),
-        exp_weight_edf=P('fsdp', None, 'tp'),
-        exp_weight_efd=P('fsdp', 'tp', None),
+        act_btd=P('fsdp', sp, None if is_sampling else 'tp'),  # pyrefly: ignore[bad-argument-type]
+        act_btf=P('fsdp', sp, 'tp'),  # pyrefly: ignore[bad-argument-type]
+        act_btnh=P('fsdp', sp, 'tp', None),  # pyrefly: ignore[bad-argument-type]
+        score_weight_d1=P(fsdp, None),  # pyrefly: ignore[bad-argument-type]
+        exp_weight_edf=P('fsdp', None, 'tp'),  # pyrefly: ignore[bad-argument-type]
+        exp_weight_efd=P('fsdp', 'tp', None),  # pyrefly: ignore[bad-argument-type]
     )
 
 
@@ -440,7 +440,7 @@ def __init__(
   ):
     self.w = nnx.Param(
         nnx.initializers.ones_init()(
-            rngs.params(), dim, param_dtype
+            rngs.params(), dim, param_dtype  # pyrefly: ignore[bad-argument-type]
         ),  # pyrefly: ignore[bad-argument-type]
         sharding=shd_config.rms_norm_weight,
     )
@@ -533,13 +533,13 @@ def block(
     value_proj = self.v_proj(x)
 
     query_proj = shard(
-        query_proj, self.shd_config.act_btnh
+        query_proj, self.shd_config.act_btnh  # pyrefly: ignore[bad-argument-type]
     )  # pyrefly: ignore[bad-argument-type]
     key_proj = shard(
-        key_proj, self.shd_config.act_btnh
+        key_proj, self.shd_config.act_btnh  # pyrefly: ignore[bad-argument-type]
     )  # pyrefly: ignore[bad-argument-type]
     value_proj = shard(
-        value_proj, self.shd_config.act_btnh
+        value_proj, self.shd_config.act_btnh  # pyrefly: ignore[bad-argument-type]
     )  # pyrefly: ignore[bad-argument-type]
 
     query_proj = apply_rope(
@@ -687,7 +687,7 @@ def sharded_splash_attn(kernel, q_block, k_block, v_block):
 
     outputs = self.o_proj(qkv)
     outputs = shard(
-        outputs, self.shd_config.act_btd
+        outputs, self.shd_config.act_btd  # pyrefly: ignore[bad-argument-type]
     )  # pyrefly: ignore[bad-argument-type]
 
     if cache is not None:
@@ -804,7 +804,7 @@ def __call__(self, x, use_megablox=True):
     if not use_megablox or (mesh.empty or jax.devices()[0].platform == 'cpu'):
       dispatch_mask = jax.nn.one_hot(
           routing_idx,
-          num_classes=self.num_experts,
+          num_classes=self.num_experts,  # pyrefly: ignore[bad-argument-type]
           dtype=self.dtype,  # pyrefly: ignore[bad-argument-type]
       )  # [B, T, K, E]
       dispatch_mask = jnp.swapaxes(dispatch_mask, -1, -2)  # [B, T, E, K]
@@ -879,12 +879,12 @@ def sharded_megablox_moe(inputs, weights, indices, gate_w, up_w, down_w):
         ep_shard_idx = 0
 
       num_local_experts = (
-          self.num_experts // num_ep
+          self.num_experts // num_ep  # pyrefly: ignore[unsupported-operation]
       )  # pyrefly: ignore[unsupported-operation]
 
       flat_repeated_inputs = jnp.repeat(
           inputs.reshape(B * T, D_global),
-          self.experts_per_tok,
+          self.experts_per_tok,  # pyrefly: ignore[bad-argument-type]
           axis=0,  # pyrefly: ignore[bad-argument-type]
       )
       flat_selected_indices = indices.reshape(-1)
@@ -918,7 +918,7 @@ def sharded_megablox_moe(inputs, weights, indices, gate_w, up_w, down_w):
         local_output_offsets = global_out_offsets[ep_shard_idx]
 
         output_buffer_size = (
-            min(self.experts_per_tok, num_local_experts)
+            min(self.experts_per_tok, num_local_experts)  # pyrefly: ignore[bad-specialization]
             * B
             * T
             * num_ep  # pyrefly: ignore[bad-specialization]
@@ -1095,7 +1095,7 @@ def block(
   ) -> jaxtyping.Array:
     activations = nnx.silu(self.gate_proj(x)) * self.up_proj(x)
     activations = shard(
-        activations, self.shd_config.act_btf
+        activations, self.shd_config.act_btf  # pyrefly: ignore[bad-argument-type]
     )  # pyrefly: ignore[bad-argument-type]
     outputs = self.down_proj(activations)
     return outputs
diff --git a/tunix/rl/agentic/agentic_grpo_learner.py b/tunix/rl/agentic/agentic_grpo_learner.py
index 00700bca9..431408859 100644
--- a/tunix/rl/agentic/agentic_grpo_learner.py
+++ b/tunix/rl/agentic/agentic_grpo_learner.py
@@ -255,7 +255,7 @@ def __init__(
             "algo_config": self.algo_config,
         }
     )
-    self.rl_engine.actor_trainer.with_rl_metrics_to_log({
+    self.rl_engine.actor_trainer.with_rl_metrics_to_log({  # pyrefly: ignore[bad-argument-type]
         "kl": common.mean_of_means,
         "entropy": common.mean_of_means,
         "reduced_pg_loss": common.mean_of_means,
@@ -455,7 +455,7 @@ def _compute_packed_logps(self, example: TrainExample) -> TrainExample:
         updates["old_per_token_logps"] = trainer_logps
 
     if updates:
-      example = example.replace(**updates)
+      example = example.replace(**updates)  # pyrefly: ignore[missing-attribute]
     return example
 
   def _process_results(
diff --git a/tunix/rl/agentic/agentic_rl_learner.py b/tunix/rl/agentic/agentic_rl_learner.py
index 15dcac0cd..e91a3e60b 100644
--- a/tunix/rl/agentic/agentic_rl_learner.py
+++ b/tunix/rl/agentic/agentic_rl_learner.py
@@ -845,7 +845,7 @@ def _to_train_examples(raw_gen):
       # independent of any micro-batch/streaming granularity.
       train_data_gen = rl_utils.pack_sequences(
           train_data_gen,
-          self._training_config.max_seq_token_per_tpu,
+          self._training_config.max_seq_token_per_tpu,  # pyrefly: ignore[bad-argument-type]
           sequences_per_update=mini_batch_size * self._num_generations(),
           pack_size=pack_size,
           max_segments_per_packed_row=getattr(
@@ -905,7 +905,7 @@ def _to_train_examples(raw_gen):
             jnp.array(float((seg == 0).sum())), jnp.array(float(seg.size))
         )
         self.rl_engine.buffer_metrics_async(
-            {
+            {  # pyrefly: ignore[bad-argument-type]
                 "packing/dummy_ratio": (
                     dummy_ratio,
                     common.global_weighted_mean,
diff --git a/tunix/rl/agentic/trajectory/trajectory_collect_engine.py b/tunix/rl/agentic/trajectory/trajectory_collect_engine.py
index 6652f863b..29ea2065f 100644
--- a/tunix/rl/agentic/trajectory/trajectory_collect_engine.py
+++ b/tunix/rl/agentic/trajectory/trajectory_collect_engine.py
@@ -236,7 +236,7 @@ async def collect(self, mode: str = "Conversation") -> Any:
       )
 
     if mode == "Trajectory":
-      self.agent.trajectory.env_time = self.env_time
+      self.agent.trajectory.env_time = self.env_time  # pyrefly: ignore[bad-assignment]
       self.agent.trajectory.reward_time = self.reward_time
       return self.agent.trajectory
     elif mode == "Steps":
diff --git a/tunix/rl/algo_core.py b/tunix/rl/algo_core.py
index 2aeeb5855..dd3a51a9f 100644
--- a/tunix/rl/algo_core.py
+++ b/tunix/rl/algo_core.py
@@ -246,7 +246,7 @@ def ppo_policy_loss_fn(
   }
 
   if return_entropy:
-    unreduced_entropy = jnp.sum(token_entropy * completion_mask)
+    unreduced_entropy = jnp.sum(token_entropy * completion_mask)  # pyrefly: ignore[unbound-name]
     unreduced_policy_loss = (
         unreduced_policy_loss - entropy_coef * unreduced_entropy
     )
@@ -458,10 +458,10 @@ def grpo_loss_fn(
       # scatter each token its own segment's mean via take_along_axis. Padding
       # (segment 0, mask 0) yields 0 and is masked out downstream.
       per_seg_sum = common.segmented_sum(
-          seq_importance_ratio * completion_mask, segment_ids, num_segments
+          seq_importance_ratio * completion_mask, segment_ids, num_segments  # pyrefly: ignore[bad-argument-type]
       )
       per_seg_count = common.segmented_count(
-          segment_ids, num_segments, mask=completion_mask
+          segment_ids, num_segments, mask=completion_mask  # pyrefly: ignore[bad-argument-type]
       )
       per_seg_mean = per_seg_sum / jnp.clip(per_seg_count, min=1.0)
       seq_mean_ratio = jnp.take_along_axis(
@@ -617,7 +617,7 @@ def grpo_loss_fn(
     aux["kl_loss"] = kl_loss  # pyrefly: ignore[bad-assignment]
   if beta is not None and beta != 0.0:
     total_loss = sft_utils.WeightedMetric(
-        unreduced_pg_loss.unreduced_sum + beta * kl_loss.unreduced_sum,
+        unreduced_pg_loss.unreduced_sum + beta * kl_loss.unreduced_sum,  # pyrefly: ignore[unbound-name]
         unreduced_pg_loss.denominator,
         eps=unreduced_pg_loss.eps,
         min_denom=unreduced_pg_loss.min_denom,
@@ -632,7 +632,7 @@ def grpo_loss_fn(
   )
   aux["entropy"] = entropy_loss
 
-  return sft_utils.LossOutput(primary_loss=total_loss, aux_metrics=aux)
+  return sft_utils.LossOutput(primary_loss=total_loss, aux_metrics=aux)  # pyrefly: ignore[bad-argument-type]
 
 
 @function_registry.register_advantage_estimator("grpo")
diff --git a/tunix/rl/ppo/ppo_learner.py b/tunix/rl/ppo/ppo_learner.py
index 0646532c5..e4db27fce 100644
--- a/tunix/rl/ppo/ppo_learner.py
+++ b/tunix/rl/ppo/ppo_learner.py
@@ -235,7 +235,7 @@ def __init__(
         actor_rl_metrics_to_log  # pyrefly: ignore[bad-argument-type]
     )
 
-    self.rl_engine.critic_trainer.with_rl_metrics_to_log({
+    self.rl_engine.critic_trainer.with_rl_metrics_to_log({  # pyrefly: ignore[bad-argument-type]
         "vpred_mean": common.mean_of_means,
         "vf_clipfrac": common.mean_of_means,
     })
diff --git a/tunix/rl/rl_cluster.py b/tunix/rl/rl_cluster.py
index 14bd15a96..32c6175b1 100644
--- a/tunix/rl/rl_cluster.py
+++ b/tunix/rl/rl_cluster.py
@@ -983,7 +983,7 @@ def generate(
     logits = None
     if outputs[0].logits is not None:
       logits = list(
-          itertools.chain.from_iterable(out.logits for out in outputs)
+          itertools.chain.from_iterable(out.logits for out in outputs)  # pyrefly: ignore[bad-argument-type]
       )
 
     return base_rollout.RolloutOutput(
diff --git a/tunix/rl/rollout/vllm_rollout.py b/tunix/rl/rollout/vllm_rollout.py
index 8f9cd3f5d..6fa0902ce 100644
--- a/tunix/rl/rollout/vllm_rollout.py
+++ b/tunix/rl/rollout/vllm_rollout.py
@@ -53,7 +53,7 @@ def __init__(
             mapping_config=mapping_config,
             return_logprobs=rollout_config.return_logprobs,
             init_with_random_weights=rollout_config.rollout_vllm_init_with_random_weights,
-            tpu_backend_type=rollout_config.rollout_vllm_tpu_backend_type,
+            tpu_backend_type=rollout_config.rollout_vllm_tpu_backend_type,  # pyrefly: ignore[bad-argument-type]
             additional_config=rollout_config.rollout_vllm_additional_config,
             enable_dp_attention=rollout_config.rollout_vllm_enable_dp_attention,
             hbm_utilization=rollout_config.rollout_vllm_hbm_utilization,
@@ -105,7 +105,7 @@ def generate(
         temperature=rollout_config.temperature,
         top_p=rollout_config.top_p,
         top_k=rollout_config.top_k,
-        seed=rollout_config.seed,
+        seed=rollout_config.seed,  # pyrefly: ignore[bad-argument-type]
         echo=False,
         pad_output=True,
         **kwargs,
@@ -114,9 +114,9 @@ def generate(
     return base_rollout.RolloutOutput(
         text=self.output.text,
         logits=None,
-        tokens=self.output.tokens,
+        tokens=self.output.tokens,  # pyrefly: ignore[bad-argument-type]
         left_padded_prompt_tokens=self.output.padded_prompt_tokens,
-        logprobs=self.output.logprobs,
+        logprobs=self.output.logprobs,  # pyrefly: ignore[bad-argument-type]
     )
 
   def get_per_token_logps(
diff --git a/tunix/rl/utils.py b/tunix/rl/utils.py
index c4e1359f4..9dbb4245b 100644
--- a/tunix/rl/utils.py
+++ b/tunix/rl/utils.py
@@ -249,7 +249,7 @@ def _get_new_sharding(x):
   return params_on_memory_kind
 
 def create_critic_model(
-    actor_model: nnx.Module, seed: int = 0, rngs: nnx.Rngs = None, lm_head_to_replace: str = "lm_head"
+    actor_model: nnx.Module, seed: int = 0, rngs: nnx.Rngs = None, lm_head_to_replace: str = "lm_head"  # pyrefly: ignore[bad-function-definition]
 ) -> nnx.Module:
   """Creates a critic model from an actor model."""
 
@@ -274,7 +274,7 @@ def create_critic_model(
 
   # If Qwix is active for the model, also assign qwix_path for the new head
   if hasattr(critic_model, "qwix_path"):
-    new_head.qwix_path = getattr(lm_head, "qwix_path", (lm_head_to_replace,))
+    new_head.qwix_path = getattr(lm_head, "qwix_path", (lm_head_to_replace,))  # pyrefly: ignore[missing-attribute]
   setattr(
       critic_model,
       lm_head_to_replace,
@@ -294,8 +294,8 @@ def __init__(self, transformer: nnx.Module, rngs: nnx.Rngs):
     """
     if hasattr(transformer, 'embed_dim'):
       embed_dim = transformer.embed_dim
-    elif hasattr(transformer.config, 'embed_dim'):
-      embed_dim = transformer.config.embed_dim
+    elif hasattr(transformer.config, 'embed_dim'):  # pyrefly: ignore[missing-attribute]
+      embed_dim = transformer.config.embed_dim  # pyrefly: ignore[missing-attribute]
     else:
       raise ValueError("Could not determine embed dim for the transformer.")
 
@@ -306,7 +306,7 @@ def __init__(self, transformer: nnx.Module, rngs: nnx.Rngs):
         use_bias=False,
         kernel_init=nnx.with_partitioning(
             nnx.initializers.normal(),
-            transformer.config.shd_config.score_weight_d1,
+            transformer.config.shd_config.score_weight_d1,  # pyrefly: ignore[missing-attribute]
         ),
         rngs=rngs,
     )
@@ -379,12 +379,12 @@ def unpad_train_example(example: common.TrainExample) -> list[dict[str, Any]]:
         "completion_mask": c_mask[i, :c_len],
         "advantages": adv[i, :c_len] if adv_is_per_token else adv[i],
         "adv_is_per_token": adv_is_per_token,
-        "ref_per_token_logps": ref_logps[i, :c_len] if has_ref else None,
-        "old_per_token_logps": old_logps[i, :c_len] if has_old else None,
-        "returns": returns_np[i, :c_len] if has_returns else None,
-        "old_values": old_values_np[i, :c_len] if has_old_values else None,
+        "ref_per_token_logps": ref_logps[i, :c_len] if has_ref else None,  # pyrefly: ignore[unbound-name]
+        "old_per_token_logps": old_logps[i, :c_len] if has_old else None,  # pyrefly: ignore[unbound-name]
+        "returns": returns_np[i, :c_len] if has_returns else None,  # pyrefly: ignore[unbound-name]
+        "old_values": old_values_np[i, :c_len] if has_old_values else None,  # pyrefly: ignore[unbound-name]
         "policy_version": (
-            policy_version_np[i : i + 1] if has_policy_version else None
+            policy_version_np[i : i + 1] if has_policy_version else None  # pyrefly: ignore[unbound-name]
         ),
     }
     res.append(item)
@@ -672,7 +672,7 @@ def _mark(merged, is_update):
     # once. Set here (not per bin) so every emitted chunk carries it.
     kwargs = dict(is_update_step=jnp.array([is_update], dtype=jnp.bool_))
     if hasattr(merged, "num_segments"):
-      kwargs["num_segments"] = effective_max_segments + 1
+      kwargs["num_segments"] = effective_max_segments + 1  # pyrefly: ignore[bad-assignment]
     return [merged.replace(**kwargs)]
 
   # See the docstring: buffer sequences, emit a chunk once it holds a chunk's
diff --git a/tunix/utils/topology.py b/tunix/utils/topology.py
index 3a3016346..84396637f 100644
--- a/tunix/utils/topology.py
+++ b/tunix/utils/topology.py
@@ -560,7 +560,7 @@ def supported_topology_shapes_for_chip_count(
     if chip_rank == 2:
       return [shape[:2] for shape in supported_shapes]
     if chip_rank == 3:
-      return supported_shapes
+      return supported_shapes  # pyrefly: ignore[bad-return]
     return []
 
   if chip_rank != 3:
@@ -580,11 +580,11 @@ def supported_topology_shapes_for_chip_count(
       required_chips, parsed_available_shape
   )
   if single_host_shapes:
-    return single_host_shapes
+    return single_host_shapes  # pyrefly: ignore[bad-return]
 
   cube_shapes = _supported_fish_cube_shapes(required_chips, parsed_available_shape)
   if cube_shapes:
-    return cube_shapes
+    return cube_shapes  # pyrefly: ignore[bad-return]
 
   return []