Skip to content

[Trellis]Qwen3.5-35B; Weight Conversion and Raiden Weight Sync Integration - #2083

Open
YixuanWang-99 wants to merge 29 commits into
mainfrom
yixuann-debug-raiden
Open

[Trellis]Qwen3.5-35B; Weight Conversion and Raiden Weight Sync Integration#2083
YixuanWang-99 wants to merge 29 commits into
mainfrom
yixuann-debug-raiden

Conversation

@YixuanWang-99

@YixuanWang-99 YixuanWang-99 commented Sep 2, 2026

Copy link
Copy Markdown
Collaborator

Description

Integrates Raiden weight synchronization with MaxText trainer and vLLM rollout workers for distributed RL workloads (GRPO / Trellis). Refactors destination-side weight sync into a shared mixin, handles multi-host FFI synchronization, manages cache lifecycles, and optimizes host memory usage.

Companion MaxText PR: AI-Hypercomputer/maxtext#5089


Key Changes

  • MaxText Trainer Integration (run_trainer_node.py, maxtext_utils.py): Adds trainer_backend="maxtext" support, configures mesh topologies (FSDP, TP, EP), enables weight conversion, and aligns trainer TP with rollout TP.
  • RaidenDestinationWeightSyncMixin Refactoring (raiden_weight_sync_mixin.py): Extracts a reusable mixin consolidating get_weight_sync_metadata, bind_weight_sync, sync_weights, and release_weight_sync across sampler adapters (inprocess_vllm_sampler_adapter, vanilla_sampler_adapter).
  • Multi-Host FFI & Step-0 Sync (raiden_handler.py, raiden_synchronizer.py, rl_program.py): Supports comma-separated multi-host control_plane_rpc_address registrations, forces weight sync at step 0 to prevent uninitialized generation, and adds tensor/element count verification to checksum diagnostics.
  • Cache Lifecycle & Memory Optimization (raiden_weight_sync_delegate.py): Adds prefix cache reset and KV cache reinitialization; adds --enable_prefix_caching flag (default false) to prevent recurrent state desync in hybrid models; releases host arrays post-sync.
  • 128-Lane MoE Chunking & Launch Script (tunix.generate.utils, launch_raiden.sh): Aligns MoE interleaving to 128-element lane chunks for TPU GMM v2 kernels, and provides unified cluster deployment and triage automation.

Testing

  • Unit Tests: Updated and verified raiden_handler_test.py, raiden_synchronizer_test.py, and weight_sync_test.py.
  • E2E Verification: Validated full 2-step distributed RL loop on Cloud TPU v5p with MaxText trainer and vLLM rollout workers, logs

Checklist

  • Unit tests added and verified.
  • Verified compatibility with companion MaxText PR.
  • Verified on multi-node TPU cluster.

A9isha and others added 11 commits September 3, 2026 19:18
The distributed runtime imports discovery_service_pb2 and
discovery_service_pb2_grpc, but only the .proto was tracked. That works in an
editable install, where the generated files sit in the working tree, and fails
anywhere the package is installed from a GitHub archive -- an archive contains
only tracked files, so the stubs are absent and every worker process dies at
import:

  ImportError: cannot import name 'discovery_service_pb2' from
  'tunix.experimental.distributed.runtime.discovery' (unknown location)

Observed on all three worker pods of a GKE run. Nothing in the packaging
generates them: there is no protoc step in pyproject.toml, and the .proto is not
shipped in the wheel either, so generating at install time is not an option
without also packaging the source.

Generated with grpcio-tools 1.81.1 / protobuf 6.33.6, which stamps gencode
6.33.5. The gencode version must not exceed the protobuf runtime wherever these
execute, and the post-training image ships protobuf 6.33.6; stubs built with a
7.x toolchain load fine locally but abort there with

  VersionError: Detected incompatible Protobuf Gencode/Runtime versions ...
  gencode 7.35.1 runtime 6.33.6

Regenerate with a protobuf 6.x toolchain, from the repo root so the imports are
fully qualified.
CPU_MACHINE, TPU_SLICE and GCS_SCRATCH_LOCATION were hardcoded to what
trellis-demo-0810 happened to have, so the launcher only ran on that
cluster. mlperf-v5p has no n2-standard-64 pool and no 2x2x2 topology, and
Pathways' default compilation-cache bucket is not writable from this
project -- which kills the compilation service rather than degrading it.
All three are now ${VAR:-<previous default>}, so existing invocations are
unchanged.

The jobsets gain HF_TOKEN from a secretKeyRef. tunix/oss/utils.py calls
hf.login() whenever HF_TOKEN is unset, even for a public model, which in a
pod polls for a device code until it fails. Marked optional: true so pods
without the secret still start.

The secret name is still hardcoded to anisha-hf-token and should be
parameterised before this goes anywhere beyond the current cluster.
maxtext_engine.py's release_weight_sync() docstring already claimed to
release staged weight buffers after transfer completion, but the
implementation only logged metrics -- self.arrays (the host-staged
copy) kept lingering for the entire idle window between rounds with
nothing to actually free it.

This alone does not shrink the native transport's own hold lifetime
(BindWeights only releases the previous round's hold atomically with
acquiring the next one -- see the chunking fix in maxtext, commit
7f90d9c53, for the actual peak-memory fix), but it keeps our own
Python-side reference from outliving its purpose once d2h() has
already copied the data into the native transport's persistent
buffer, and it makes the intent of release_weight_sync() match what
its docstring already claimed. self.names is left untouched so `bound`
keeps reporting whether bind() has ever run.

(cherry picked from commit 1fc6ff9)
Ports vllm_sampler_adapter.py from mohit/raiden-maxtext-rlvllm, which
drives tpu-inference's RLVllmSampler (SAMPLER=vllm) instead of running
a vLLM engine in-process. Imports are repointed at the weight_sync
package, since the orchestrator module this branch was written against
has since been split out.

That repointing is also why dict_to_metadata comes along here: it lives
in orchestrator/weight_sync.py on the source branch and had no
counterpart under weight_sync/, so the adapter resolved its module but
not the symbol, and only failed once weight sync round 0 actually
called it. The two paths need it for different reasons -- an in-process
destination returns WorkUnitMetadata directly, but this one binds
Raiden inside a separate EngineCore process, so tpu-inference flattens
the same content to plain dicts to cross that boundary. Rebuilding the
dataclasses here lets both present one type to manifest preflight.
Four fixes, each needed before the rollout node would start under
--sampler=vllm. All were found by running it; none are speculative.

Import tpu_sync's native extension first. Loading it after vLLM/torch
pull in their own copy aborts the process with "free(): invalid
pointer" inside tpu_inference.rl.raiden_worker_sync, which imports the
same module. Bisecting the three sampler modes showed the crash is
independent of weight sync, so the guarded import goes above every
other import and is load-bearing rather than stylistic.

Resolve --tensor_parallel_size after parsing instead of defaulting it
to jax.device_count(). Evaluating that default opens the TPU in the
parent while the parser is still being built; the vllm path then hands
those same chips to a separate EngineCore process, which either cannot
reopen them ("Device or resource busy") or inherits them half-owned via
fork and hangs on its first compile. The launcher already says which
chips we own, so read that and only fall back to JAX when it is unset.

Set NEW_MODEL_DESIGN for the MaxText branch. MaxText's inference
vllm.yml declares a five-axis mesh, but tpu-inference only builds one
under that env; its default 2D ('data','model') mesh fails model
loading with "Resource axis: attn_dp ... is not found in mesh". Scoped
to this branch so the HF path keeps the 2D mesh its MoE kernel wants.

Thread MAXTEXT_MODEL_NAME from one launcher variable into both the
trainer and the rollout. Defaulting the two sides separately is how
they drift, and drift here is not a clean failure: Raiden pairs tensors
by exact name, so a MaxText trainer against a non-MaxText rollout
matches zero names.
The failures ride on WeightSyncError structurally, but the default
renderer prints only the outer message, so a mismatch reported
"preflight failed" without naming anything. Problems are name-sorted,
so a whole-convention mismatch fills the head of the list with one side
only; sample both manifests explicitly, since that is what identifies
the convention gap and the problem lines alone cannot show it.
Nothing on this path records the completion, so a run that generates
nothing but newlines reports the same "rollouts=N" summary as a real
one -- which is exactly how a broken Mode 2 rollout was mistaken for a
working one. Logs token counts and a text prefix per trajectory.

Revert once the vllm sampler path is trusted; this is a debugging aid,
not a feature.
train_stage only syncs after a step, so step-0 trajectories came from
whatever the rollout worker built for itself. On the MaxText-in-vLLM path
that is a randomly initialized model: MaxTextForCausalLM is constructed
with an empty load_parameters_path, so from_pretrained skips the Orbax
restore. The resulting degenerate rollouts were still scored and trained
on, so the run looked healthy while learning from noise.

StandardRLProgram now pushes the trainer's starting weights out before
dispatching anything, retrying until the vLLM EngineCore subprocess comes
up (~20s locally) and failing loudly at the deadline rather than rolling
out from uninitialized weights.

That sync needs the engine, which RLVllmSampler only builds lazily on the
first sample(), so VllmSamplerAdapter starts it on demand from the two
weight-sync entry points. Without that the round finds no worker, reports
an empty destination manifest, and the run deadlocks with the engine
waiting for a sample that dispatch is waiting on the sync to allow.
The source and destination grand totals are only comparable if both sides
bound the same tensors, which the checksum dict left implicit. Emitting
__tensor_count__ and __element_count__ alongside makes that checkable
rather than assumed; tpu-inference's RaidenWorkerSync.checksums() reports
the same keys.
…istributed configs

- In RaidenSynchronizer, flatten state with incremental device-to-cpu transfers, GC, and malloc_trim to minimize peak host memory
- Support flat_state attribute in transfer_state_directly for nnx state updates
- Make is_bounded check idempotent in sampler adapters
- Increase PhaseTimeouts bind and metadata limits to 180s
- Add CLI options and auto MoE padding calculation for MaxText trainer/rollout nodes
…ode.py

- Set vllm.use_weight_converter and vllm.rollout_backend in MaxText config argv in run_trainer_node.py
…depool toleration

- Add gke-nodepool tolerations in jobset.cpu.yaml
- Replace TP arguments with --tensor_parallel_size in run_rollout_node.py
- Use thread pool executor for submit_worker in DistributedRLEngine when loop is running
- Strip .value suffix from variable names in VllmSamplerAdapter, flatten_weights, and WorkUnitMetadata
- Support host_stage mode in RaidenSynchronizer under proxy backend
…th Raiden sync

- Provide one-command launcher, monitor, triage, and teardown for MaxText + Tunix + vLLM + Raiden workloads
- Support model presets (qwen3-0.6b, qwen3.5-35b, qwen3-1.7b) with verified slice topologies and checkpoints
- Include automated log triaging for weight conversion and hardware errors
@YixuanWang-99 YixuanWang-99 changed the title [DO NOT REVIEW][Trellis]Weight conversion that compatible with raiden [Trellis]Weight conversion that compatible with raiden Sep 3, 2026
…mands in launch_raiden.sh

- Truncate default RUN_ID username prefix to stay within Kubernetes 63-character label value limit for coordinator labels
- Add explicit length check rejecting TRAINER_ID > 26 characters
- Add start-trainer and start-rollout subcommands
- Handle render and dry-run commands cleanly
- Wait for JobSet deletion during workload teardown
…replica unit ID

- Thread --rollout_replicas from launch_raiden.sh to run_gsm8k_dist_grpo.py
- Wait for all rollout worker replicas in cluster.wait_for_workers()
- Set unit=server_id in VllmSamplerAdapter.get_raiden_metadata() to disambiguate replicas
- Eagerly drain completed responses in DistributedRLEngine.poll_rollouts() with non-blocking poll loop
…date test assertions

- Dynamically register custom call targets via libraiden_ffi_bridge.so in RaidenSynchronizer
- Parameterize HF_TOKEN_SECRET_NAME in yaml_generator and jobset templates with maxRestarts=3
- Deduplicate dict_to_metadata in weight_sync.py
- Update checksums test assertions in raiden_synchronizer_test.py for tensor and element counts
- Clean up orphaned batch jobs and default SYNC_CODE=true in launch_raiden.sh
…oyment configuration

- Support comma-separated multi-axis sharding specifications and monkey-patch TPUWorker in raiden_synchronizer
- Parameterize k8s namespace and improve status/log triage in launch_raiden.sh and deployment YAMLs
- Respect DISABLE_CHECKPOINTING environment variable in run_trainer_node and rl_program
- Add base_num_kv_heads configuration parameter in maxtext_utils
- Lower bind and metadata timeouts and improve shard metadata logging in weight_sync_coordinator
…nable MoE weight prefusion, and support multi-host FFI synchronization
…efault: false) to prevent recurrent state desync in hybrid models
…lifecycle management

- Extract RaidenDestinationWeightSyncMixin for rollout sampler adapters
- Support prefix cache reset and KV cache deletion/reinitialization in RaidenWeightSyncDelegate
- Release host arrays after weight sync in PeftTrainer
- Guard patch_raiden_worker_sync when JAX_PLATFORMS=cpu
- Define TPU_V5P_SUBCORE_LANE_SIZE constant in tunix.generate.utils
- Remove stale launch_raiden.sh symlink
@YixuanWang-99 YixuanWang-99 changed the title [Trellis]Weight conversion that compatible with raiden [Trellis]Weight Conversion and Raiden Weight Sync Integration Sep 8, 2026
@YixuanWang-99 YixuanWang-99 changed the title [Trellis]Weight Conversion and Raiden Weight Sync Integration [Trellis]Qwen3.5-35B; Weight Conversion and Raiden Weight Sync Integration Sep 8, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why do we want to add these? these should be auto generated on setup.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

+1 please revert

completion = action.action if hasattr(action, "action") else str(action)
reward, info = gsm8k_env_reward(self.task, action)
info["correct"] = bool(info["answer_correct"])
prompt_text = self.task.get("question", "") or self.task.get("prompts", "")

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is debugging logs, we should push all these. especially at info level. For this and other changes, can we do a thorough review to make sure we don't clutter the logs

return self._trainer.prepare_weight_sync(**kwargs)

def save_checkpoint(self, metadata: Any = None, **kwargs) -> None:
if os.environ.get("DISABLE_CHECKPOINTING", "false").lower() in ("true", "1"):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

instead of os env checking can we pass through the config?


async def poll_rollouts(
self, timeout_s: float = remote_execution.LONG_POLL_TIMEOUT_S
self, timeout_s: float = 1.0

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why are we changing these values? Also in a few other locations.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

timeout of 1s doesn't make any sense, please revert

if res is None or isinstance(res, Exception):
return []
worker_items = [res]
while True:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why do we need this second while loop? the polling automatically checks in a loop too.

"num_microbatches": num_microbatches,
},
)
if os.environ.get("DISABLE_CHECKPOINTING", "false").lower() in ("true", "1"):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ditto, let's pass this properly


async def poll_rollouts(
self, timeout_s: float = remote_execution.LONG_POLL_TIMEOUT_S
self, timeout_s: float = 1.0

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ditto.

async def _sync_initial_weights(self) -> None:
"""Pushes the trainer's starting weights out before the first dispatch.

`train_stage` only syncs *after* a step, so whatever the rollout worker

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

shouldn't this change be pushed to MaxText-in-vLLM instead so uit doesn't load from an empty checkpoint? that seems like a better level of abstraction and modularity.

help="Weight sync mode (none, fallback, or raiden).",
)
return parser.parse_args(argv)
parser.add_argument(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can these improvement be split into their own cl? doesn't seem relevant for the fix iiuc.

" step 0)."
),
)
parser.add_argument(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ditto on seperate cl

@tianshub tianshub left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please split this PR in the smaller ones and make sure each one contains a well focused change

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

+1 please revert

#!/bin/bash
# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this banner should still be there

@@ -1,107 +1,889 @@
#!/bin/bash

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why do we need all these changes in the launcher script? it was runnable previously no?


async def poll_rollouts(
self, timeout_s: float = remote_execution.LONG_POLL_TIMEOUT_S
self, timeout_s: float = 1.0

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

timeout of 1s doesn't make any sense, please revert

)


def _submit_worker(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is any of the changes in this file related to weight sync? also the asyncio loop should already be running when we start invoke worker endpoints, why do we need to handle this here?

self.on_step_end(current_step, step_result)
self._step += 1

async def _sync_initial_weights(self) -> None:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm pretty sure I've seen this function somewhere in another PR, do you need to rebase and merge?

self.max_response_length = request.generation_kwargs.get(
"max_response_length"
)
self._accumulated_token_ids: List[int] = []

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why do we need these changes?

" transfer_parallelism, not both"
)
worker_rpc_client = None
if name_resolver is not None:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why is this if condition removed? weight sync is not a mandatory arg for trainer

- Point PRESET_DEFAULT_IMAGE to yixuann-dev-0908-ffi for qwen35-35b-moe-raiden
- Add flag aliases for --model-name/--model_name and --queue-name/--queue_name
@YixuanWang-99

Copy link
Copy Markdown
Collaborator Author

Thanks everyone for the detailed review feedback!

Per the feedback on PR size and scope, this PR is being split into a stacked chain of focused, reviewable PRs targeting main:

Stacked PR Chain:

  1. [Raiden Weight Sync 1/6] Support multi-host discovery and multi-axis sharding in RaidenHandler #2146 (T1): Multi-host discovery & multi-axis sharding in RaidenHandler (raiden_handler.py + tests)
    • Restores the name_resolver is not None guard (@tianshub).
  2. [Raiden Weight Sync 2/6] FFI transport selection, preflight mismatch check, and host-stage normalization #2147 (T2a): FFI transport selection, preflight mismatch check, and host-stage normalization (raiden_synchronizer.py, weight_sync_coordinator.py, weight_sync.py)
    • Fixes the Step 0 blocker where FFI was never enabled end-to-end between Pathways trainer and rollout. Adds preflight validation in WeightSyncCoordinator to reject mismatched transport modes before transfer.
    • Pairs with MaxText #5171 (M4).
  3. [Raiden Weight Sync 3/6] Extract RaidenDestinationWeightSyncMixin for rollout samplers #2148 (T3): Extract RaidenDestinationWeightSyncMixin for rollout samplers
    • Clean net −65 line refactoring extracting common destination lifecycle methods across inprocess_vllm, vanilla, and external vLLM adapters.
  4. [Raiden Weight Sync 4/6] Manage cache lifecycle and prefix caching in RaidenWeightSyncDelegate #2149 (T4): Cache lifecycle management and prefix caching
    • Prefix cache invalidation + KV state re-init upon weight transfer. Adds --enable_prefix_caching flag (default false) to avoid recurrent state desync in hybrid models.
  5. [Raiden Weight Sync 5/6] MoE 128-lane weight interleaving for TPU GMM layout #2150 (T5): MoE 128-lane weight interleaving for TPU GMM layout (tunix/generate/utils.py)
    • Matches TPU GMM kernel interleaving (TPU_V5P_SUBCORE_LANE_SIZE = 128).
    • Pairs with MaxText #5168 (M2). Includes dedicated unit tests.
  6. [Raiden Weight Sync 6/6] MaxText trainer config plumbing and MoE dimension padding #2151 (T6): MaxText trainer configuration plumbing (maxtext_utils.py)
    • Replaces ad-hoc env-var checking with explicit parameters (base_num_kv_heads, rollout_mesh_tp, prefuse_moe_weights, padded_moe_mlp_dim).
  7. [Deployment] Parameterize JobSet manifests and minimal k8s_launcher for multi-host RL #2152 (T9): Parameterize JobSet manifests and update k8s_launcher.sh
    • Independent deployment slice. Restores Apache 2.0 license banner (@tianshub).

Notes on review comments & dropped duplicates:

We will keep yixuann-debug-raiden alive as the integration branch while landing the stacked PRs.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants