[Tunix] Support scanned to unscanned weight transfer in transfer_state_directly by gagika · Pull Request #1008 · google/tunix

gagika · 2026-01-28T06:42:50Z

This PR extends the transfer_state_directly utility to support weight synchronization from scanned MaxText model (where layers are stacked in a single tensor) to unscanned MaxText + vLLM models (where layers are separate parameters).

Previously, transfer_state_directly only supported 1-to-1 mapping (Unscanned -> Unscanned). This change adds logic to detect and unroll scanned layers during the transfer process.

I have added all the necessary unit tests for my change.
I have verified that my change does not break existing code and all unit tests pass.
I have added all appropriate doc-strings/documentation.
My PR is based on the latest changes of the main branch (if unsure, rebase the code).
I have signed the Contributor License Agreement.
I have followed Contribution Guidelines.

wang2yn84 · 2026-02-03T19:19:23Z

tunix/generate/utils.py



+def _slice_scanned_param(
+    src_val: Any, tgt_val: Any, slice_idx: int, key_path: str


Can you use more detailed types instead of Any?

Done, added a few relevant types.

wang2yn84 · 2026-02-03T19:20:04Z

tunix/generate/utils.py

+def _slice_scanned_param(
+    src_val: Any, tgt_val: Any, slice_idx: int, key_path: str
+) -> Any:
+  """Slices a scanned parameter dynamically detecting the scan axis."""


Can you put a more detailed doc string? And maybe also include the input output descriptions?

wang2yn84 · 2026-02-03T19:22:27Z

tunix/generate/utils.py

-      src: Any, tgt_spec: Any, path: str = ''
-  ) -> Tuple[Any, Any]:
-    # Stop recursion if we hit a leaf (non-dict)
+  # Helper: Intersect Trees (Handle KVCache/RNG mismatches and Scanned Layers)


Can you fold it into the docstring?

wang2yn84 · 2026-02-03T19:22:38Z

tunix/generate/utils.py

-  ) -> Tuple[Any, Any]:
-    # Stop recursion if we hit a leaf (non-dict)
+  # Helper: Intersect Trees (Handle KVCache/RNG mismatches and Scanned Layers)
+  def intersect_trees(src: Any, tgt_spec: Any) -> Tuple[Any, Any]:


I understand it's there before your PR, but can you still add the detailed types?

Added more specific typing here.

wang2yn84 · 2026-02-03T19:23:11Z

tunix/generate/utils.py

-    # Stop recursion if we hit a leaf (non-dict)
+  # Helper: Intersect Trees (Handle KVCache/RNG mismatches and Scanned Layers)
+  def intersect_trees(src: Any, tgt_spec: Any) -> Tuple[Any, Any]:
+    """Optimized intersection using flat dictionary traversal."""


wang2yn84 · 2026-02-03T19:26:35Z

tunix/generate/utils.py

+  try:
+    return src_val[slice_idx]
+
+  except (IndexError, TypeError):


Add more debugging information in case of this?

wang2yn84 · 2026-02-03T19:28:50Z

tunix/generate/utils.py

+        candidate_b.pop(match_index)
+        candidate_b = tuple(candidate_b)
+
+        if candidate_b in src_flat:


Duplicate code as candidate a, consider make it simpler?

Simplified the code.

…e_directly Dynamically detecting scan dim + path caching. adding explicit cleanup. improving docstrings and typing

-- 24224a5 by Gagik Amirkhanyan <agagik@google.com>: [Tunix] Support scanned to unscanned weight transfer in transfer_state_directly Dynamically detecting scan dim + path caching. adding explicit cleanup. improving docstrings and typing COPYBARA_INTEGRATE_REVIEW=#1008 from google:agagik-scan 24224a5 PiperOrigin-RevId: 869308029

gagika temporarily deployed to testing January 28, 2026 06:43 — with GitHub Actions Inactive

NicoGrande force-pushed the agagik-scan branch from 2efc8f5 to 93b4730 Compare January 29, 2026 21:56

NicoGrande temporarily deployed to testing January 29, 2026 21:56 — with GitHub Actions Inactive

NicoGrande had a problem deploying to testing January 29, 2026 21:56 — with GitHub Actions Failure

NicoGrande force-pushed the agagik-scan branch from 93b4730 to bbf5b3b Compare January 30, 2026 05:04

NicoGrande had a problem deploying to testing January 30, 2026 05:05 — with GitHub Actions Failure

NicoGrande temporarily deployed to testing January 30, 2026 05:05 — with GitHub Actions Inactive

NicoGrande force-pushed the agagik-scan branch from bbf5b3b to 03cc987 Compare January 31, 2026 02:02

NicoGrande temporarily deployed to testing January 31, 2026 02:02 — with GitHub Actions Inactive

NicoGrande force-pushed the agagik-scan branch from 03cc987 to 4df97fc Compare February 2, 2026 23:53

NicoGrande had a problem deploying to testing February 2, 2026 23:53 — with GitHub Actions Error

NicoGrande force-pushed the agagik-scan branch from 4df97fc to feb485d Compare February 2, 2026 23:53

NicoGrande temporarily deployed to testing February 2, 2026 23:54 — with GitHub Actions Inactive

wang2yn84 reviewed Feb 3, 2026

View reviewed changes

NicoGrande force-pushed the agagik-scan branch from feb485d to c57986c Compare February 11, 2026 19:02

NicoGrande temporarily deployed to testing February 11, 2026 19:02 — with GitHub Actions Inactive

NicoGrande force-pushed the agagik-scan branch from c57986c to 3222c0b Compare February 11, 2026 20:33

NicoGrande temporarily deployed to testing February 11, 2026 20:34 — with GitHub Actions Inactive

NicoGrande had a problem deploying to testing February 11, 2026 20:34 — with GitHub Actions Failure

NicoGrande marked this pull request as ready for review February 11, 2026 20:35

NicoGrande requested review from abheesht17, hgao327, lc5211, sizhit2 and tianshub as code owners February 11, 2026 20:35

NicoGrande requested a review from jiangyangmu as a code owner February 11, 2026 20:35

NicoGrande force-pushed the agagik-scan branch from 3222c0b to 24224a5 Compare February 11, 2026 22:24

NicoGrande temporarily deployed to testing February 11, 2026 22:24 — with GitHub Actions Inactive

[Tunix] Support scanned to unscanned weight transfer in transfer_stat…

6947e0e

…e_directly Dynamically detecting scan dim + path caching. adding explicit cleanup. improving docstrings and typing

NicoGrande force-pushed the agagik-scan branch from 24224a5 to 6947e0e Compare February 12, 2026 18:48

NicoGrande temporarily deployed to testing February 12, 2026 18:48 — with GitHub Actions Inactive



		def _slice_scanned_param(
		src_val: Any, tgt_val: Any, slice_idx: int, key_path: str

Conversation

gagika commented Jan 28, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants