[Trellis]Weight conversion that compatible with raiden - #2083
Open
YixuanWang-99 wants to merge 19 commits into
Open
[Trellis]Weight conversion that compatible with raiden#2083YixuanWang-99 wants to merge 19 commits into
YixuanWang-99 wants to merge 19 commits into
Conversation
YixuanWang-99
requested review from
abheesht17,
hgao327,
jiangyangmu,
lc5211,
s-noghabi,
sizhit2,
tianshub and
wang2yn84
as code owners
September 2, 2026 04:09
YixuanWang-99
force-pushed
the
yixuann-debug-raiden
branch
2 times, most recently
from
September 3, 2026 19:13
3a7fe34 to
6e08b14
Compare
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
YixuanWang-99
force-pushed
the
yixuann-debug-raiden
branch
from
September 3, 2026 20:13
6e08b14 to
c46b832
Compare
…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
changed the title
…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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR integrates Raiden weight synchronization with MaxText trainer and vLLM rollout workers for distributed RL workloads (GRPO / Trellis) in Tunix.
Companion MaxText PR: AI-Hypercomputer/maxtext#5089
Key Changes
1. MaxText Trainer Integration (
run_trainer_node.py)trainer_backend="maxtext"handling to the distributed RL trainer node.--maxtext_model_name,--maxtext_ckpt_path,--maxtext_output_directory), mesh topologies (mesh_fsdp,mesh_tp,mesh_expert), and automatic MoE padded intermediate dimension calculation.vllm.use_weight_converter=True,vllm.rollout_backend="maxtext") into the MaxText engine configuration.DISABLE_CHECKPOINTINGto bypass disk I/O when iterating on weight-sync performance.2. vLLM Rollout & Sampler Adapter (
vllm_sampler_adapter.py,run_rollout_node.py)VllmSamplerAdapter: BridgesRLVllmSamplerfromtpu-inference(running in separateEngineCoresubprocesses) with Tunix'sWeightSyncDestinationinterface.WorkUnitMetadatafrom the plain dicts passed across process boundaries.--tensor_parallel_sizeafter CLI parsing without eagerly opening JAX TPU devices in the parent process, avoiding"Device or resource busy"collisions with vLLM'sEngineCore.tpu_syncnative extensions prior totorch/vllmimports to preventfree(): invalid pointercrashes..valuesuffixes from parameter names across both sampler metadata and weight manifests so that Flax/NNX variable paths align 1:1 with vLLM.3. Raiden Weight Synchronization Core (
raiden_synchronizer.py,weight_sync.py,weight_sync_coordinator.py)StandardRLProgram, guaranteeing rollouts never generate tokens from uninitialized dummy weights.checksums()to report__tensor_count__and__element_count__alongside grand totals, making tensor count parity explicitly checkable againsttpu-inference'sRaidenWorkerSync.RaidenSynchronizer.release_host_arrays()to release host-side Python array references immediately after device-to-host staging (d2h()).malloc_trimto minimize peak host RSS.host_stagemode under proxy backends, bypassing direct native FFI transport calls when host staging is managed externally.4. Orchestrator & Runtime Robustness (
distributed_rl_engine.py)_submit_worker()using aThreadPoolExecutorto dispatch synchronous worker RPCs when anasyncioevent loop is already running on the current thread, preventingRuntimeError: This event loop is already running.PhaseTimeoutsforbindandmetadatato 180s to accommodate larger model architectures.5. Unified Launch & Triage Tooling (
launch_raiden.sh)tunix/experimental/examples/math_gsm8k_dist/launch_raiden.sh:start,stop,restart,status,logs, andtriage.qwen3-0.6b,qwen3.5-35b, andqwen3-1.7b.cloud.google.com/gke-nodepooltolerations in CPU jobset definitions.discovery_service_pb2.py,discovery_service_pb2_grpc.py).Tests
E2E Test passed. (2×2×2 v5p trainer, 2×2×1 v5p rollout, MAX_STEPS=2)
Set time around 9/3 3:00pm
Trainer logs
Rollout logs
Rollout responses are sensible like:
Metrics makes sense as well:
Checklist