diff --git a/launch_raiden.sh b/launch_raiden.sh new file mode 120000 index 000000000..6b1ab7b90 --- /dev/null +++ b/launch_raiden.sh @@ -0,0 +1 @@ +/usr/local/google/home/yixuannwang/projects/tunix/tunix/experimental/examples/math_gsm8k_dist/launch_raiden.sh \ No newline at end of file diff --git a/requirements/maxtext_requirements.txt b/requirements/maxtext_requirements.txt index 38f0af3a1..bcedb329f 100644 --- a/requirements/maxtext_requirements.txt +++ b/requirements/maxtext_requirements.txt @@ -1,5 +1,5 @@ -maxtext @ git+https://github.com/AI-Hypercomputer/maxtext.git@8c2e29218c01b6287ba85e8d0dc9555561f7634c -maxtext-vllm-adapter @ git+https://github.com/AI-Hypercomputer/maxtext.git@8c2e29218c01b6287ba85e8d0dc9555561f7634c#subdirectory=src/maxtext/integration/vllm +maxtext @ git+https://github.com/AI-Hypercomputer/maxtext.git@d3178f4ef123e5515ebead1204f0c3be79f0b32d +maxtext-vllm-adapter @ git+https://github.com/AI-Hypercomputer/maxtext.git@d3178f4ef123e5515ebead1204f0c3be79f0b32d#subdirectory=src/maxtext/integration/vllm aqtp tokamax>=0.0.4 drjax>=0.1.4 diff --git a/tests/experimental/weight_sync/raiden_synchronizer_test.py b/tests/experimental/weight_sync/raiden_synchronizer_test.py index 2bbde8b44..0cd0b1b2a 100644 --- a/tests/experimental/weight_sync/raiden_synchronizer_test.py +++ b/tests/experimental/weight_sync/raiden_synchronizer_test.py @@ -222,7 +222,9 @@ def test_checksums_grand_total(self): sync = raiden_synchronizer.RaidenSynchronizer("rollout", self._state()) sums = sync.checksums() self.assertEqual(sums["__grand_total__"], 8.0 + 3.0) - self.assertLen(sums, 3) # two sampled tensors + grand total + self.assertEqual(sums["__tensor_count__"], 2) + self.assertEqual(sums["__element_count__"], 11) + self.assertLen(sums, 5) # two sampled tensors + grand total + tensor/element counts def test_work_unit_metadata_shards_and_addresses(self): sync = raiden_synchronizer.RaidenSynchronizer( diff --git a/tunix/experimental/common/datatypes.py b/tunix/experimental/common/datatypes.py index bbf98b48d..0eb203621 100644 --- a/tunix/experimental/common/datatypes.py +++ b/tunix/experimental/common/datatypes.py @@ -465,9 +465,11 @@ def _get_step_attr(step, attr): if raw_reward is None: raw_reward = getattr(traj, "reward", 0.0) try: - env_reward = float(raw_reward or 0.0) - except (ValueError, TypeError): - env_reward = 0.0 + env_reward = float(0.0 if raw_reward is None else raw_reward) + except (ValueError, TypeError) as exc: + raise ValueError( + f"Invalid reward '{raw_reward}' for request '{request_id}': must be a float." + ) from exc return cls( request_id=request_id, diff --git a/tunix/experimental/distributed/deployment/yaml_generator.py b/tunix/experimental/distributed/deployment/yaml_generator.py index 738746ede..ff6528814 100644 --- a/tunix/experimental/distributed/deployment/yaml_generator.py +++ b/tunix/experimental/distributed/deployment/yaml_generator.py @@ -89,6 +89,21 @@ def main() -> None: default="sleep infinity", help="Command to run on startup", ) + parser.add_argument( + "--hf_token_secret_name", + default=os.environ.get("HF_TOKEN_SECRET_NAME", "hf-token-secret"), + help="Kubernetes secret name containing HF_TOKEN", + ) + parser.add_argument( + "--namespace", + default=os.environ.get("K8S_NAMESPACE", "default"), + help="Kubernetes namespace (default: default)", + ) + parser.add_argument( + "--queue_name", + default=os.environ.get("KUEUE_QUEUE_NAME", ""), + help="Kueue local queue name (optional)", + ) args = parser.parse_args() @@ -164,6 +179,9 @@ def main() -> None: USER_CONTAINER_IMAGE=args.worker_container_image, USER_CONTAINER_PORT=args.worker_container_port, STARTUP_COMMAND=args.worker_startup_command, + HF_TOKEN_SECRET_NAME=args.hf_token_secret_name, + NAMESPACE=args.namespace or "default", + QUEUE_NAME=args.queue_name, ) print(content) diff --git a/tunix/experimental/distributed/deployment/yamls/jobset.cpu.yaml b/tunix/experimental/distributed/deployment/yamls/jobset.cpu.yaml index 2ab0562b3..7a4ed26a2 100644 --- a/tunix/experimental/distributed/deployment/yamls/jobset.cpu.yaml +++ b/tunix/experimental/distributed/deployment/yamls/jobset.cpu.yaml @@ -16,9 +16,10 @@ apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet metadata: name: ${JOBSET_NAME} - namespace: default + namespace: ${NAMESPACE} spec: failurePolicy: + maxRestarts: 3 restartStrategy: Recreate network: enableDNSHostnames: true @@ -40,12 +41,17 @@ spec: # ensure exclusive access to the node alpha.jobset.sigs.k8s.io/exclusive-topology: kubernetes.io/hostname spec: + priorityClassName: medium dnsPolicy: ClusterFirstWithHostNet hostNetwork: true restartPolicy: Never terminationGracePeriodSeconds: 10 nodeSelector: node.kubernetes.io/instance-type: ${CPU_MACHINE} + tolerations: + - key: "cloud.google.com/gke-nodepool" + operator: "Exists" + effect: "NoSchedule" containers: - name: ${USER_CONTAINER} image: ${USER_CONTAINER_IMAGE} @@ -83,6 +89,12 @@ spec: exit $$EXIT_CODE env: # Injects the JobSet Name + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: ${HF_TOKEN_SECRET_NAME} + key: HF_TOKEN + optional: true - name: JOBSET_NAME valueFrom: fieldRef: diff --git a/tunix/experimental/distributed/deployment/yamls/jobset.pathways.yaml b/tunix/experimental/distributed/deployment/yamls/jobset.pathways.yaml index 7dcc2a1fa..166e7423e 100644 --- a/tunix/experimental/distributed/deployment/yamls/jobset.pathways.yaml +++ b/tunix/experimental/distributed/deployment/yamls/jobset.pathways.yaml @@ -16,11 +16,12 @@ apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet metadata: name: ${JOBSET_NAME} - namespace: default + namespace: ${NAMESPACE} spec: coordinator: replicatedJob: proc failurePolicy: + maxRestarts: 3 restartStrategy: Recreate network: enableDNSHostnames: true @@ -38,6 +39,7 @@ spec: parallelism: 1 template: spec: + priorityClassName: medium dnsPolicy: ClusterFirstWithHostNet hostNetwork: true restartPolicy: Never @@ -152,6 +154,12 @@ spec: exit $$EXIT_CODE env: # Injects the JobSet Name + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: ${HF_TOKEN_SECRET_NAME} + key: HF_TOKEN + optional: true - name: JOBSET_NAME valueFrom: fieldRef: @@ -197,6 +205,7 @@ spec: # kueue.x-k8s.io/podset-slice-required-topology: cloud.google.com/gke-tpu-slice-${PODSET_SLICE_TOPOLOGY}-id # kueue.x-k8s.io/podset-slice-size: "${PODSET_SLICE_SIZE}" spec: + priorityClassName: medium dnsPolicy: ClusterFirstWithHostNet hostNetwork: true restartPolicy: OnFailure diff --git a/tunix/experimental/distributed/deployment/yamls/jobset.tpu.yaml b/tunix/experimental/distributed/deployment/yamls/jobset.tpu.yaml index dad4f44c1..609d27eed 100644 --- a/tunix/experimental/distributed/deployment/yamls/jobset.tpu.yaml +++ b/tunix/experimental/distributed/deployment/yamls/jobset.tpu.yaml @@ -16,9 +16,10 @@ apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet metadata: name: ${JOBSET_NAME} - namespace: default + namespace: ${NAMESPACE} spec: failurePolicy: + maxRestarts: 3 restartStrategy: Recreate network: enableDNSHostnames: true @@ -33,6 +34,7 @@ spec: parallelism: ${PARALLELISM} template: spec: + priorityClassName: medium dnsPolicy: ClusterFirstWithHostNet hostNetwork: true restartPolicy: OnFailure diff --git a/tunix/experimental/distributed/runtime/discovery/discovery_service_pb2.py b/tunix/experimental/distributed/runtime/discovery/discovery_service_pb2.py new file mode 100644 index 000000000..1b29a48a5 --- /dev/null +++ b/tunix/experimental/distributed/runtime/discovery/discovery_service_pb2.py @@ -0,0 +1,40 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# NO CHECKED-IN PROTOBUF GENCODE +# source: tunix/experimental/distributed/runtime/discovery/discovery_service.proto +# Protobuf Python Version: 6.33.5 +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import runtime_version as _runtime_version +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder +_runtime_version.ValidateProtobufRuntimeVersion( + _runtime_version.Domain.PUBLIC, + 6, + 33, + 5, + '', + 'tunix/experimental/distributed/runtime/discovery/discovery_service.proto' +) +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\nHtunix/experimental/distributed/runtime/discovery/discovery_service.proto\x12\x30tunix.experimental.distributed.runtime.discovery\"C\n\x0fRegisterRequest\x12\x10\n\x08hostname\x18\x01 \x01(\t\x12\x0c\n\x04port\x18\x02 \x01(\r\x12\x10\n\x08metadata\x18\x03 \x01(\x0c\"\x12\n\x10RegisterResponse2\xa6\x01\n\x10\x44iscoveryService\x12\x91\x01\n\x08Register\x12\x41.tunix.experimental.distributed.runtime.discovery.RegisterRequest\x1a\x42.tunix.experimental.distributed.runtime.discovery.RegisterResponseb\x06proto3') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'tunix.experimental.distributed.runtime.discovery.discovery_service_pb2', _globals) +if not _descriptor._USE_C_DESCRIPTORS: + DESCRIPTOR._loaded_options = None + _globals['_REGISTERREQUEST']._serialized_start=126 + _globals['_REGISTERREQUEST']._serialized_end=193 + _globals['_REGISTERRESPONSE']._serialized_start=195 + _globals['_REGISTERRESPONSE']._serialized_end=213 + _globals['_DISCOVERYSERVICE']._serialized_start=216 + _globals['_DISCOVERYSERVICE']._serialized_end=382 +# @@protoc_insertion_point(module_scope) diff --git a/tunix/experimental/distributed/runtime/discovery/discovery_service_pb2_grpc.py b/tunix/experimental/distributed/runtime/discovery/discovery_service_pb2_grpc.py new file mode 100644 index 000000000..6c9ebaf9d --- /dev/null +++ b/tunix/experimental/distributed/runtime/discovery/discovery_service_pb2_grpc.py @@ -0,0 +1,101 @@ +# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! +"""Client and server classes corresponding to protobuf-defined services.""" +import grpc +import warnings + +from tunix.experimental.distributed.runtime.discovery import discovery_service_pb2 as tunix_dot_experimental_dot_distributed_dot_runtime_dot_discovery_dot_discovery__service__pb2 + +GRPC_GENERATED_VERSION = '1.81.1' +GRPC_VERSION = grpc.__version__ +_version_not_supported = False + +try: + from grpc._utilities import first_version_is_lower + _version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION) +except ImportError: + _version_not_supported = True + +if _version_not_supported: + raise RuntimeError( + f'The grpc package installed is at version {GRPC_VERSION},' + + ' but the generated code in tunix/experimental/distributed/runtime/discovery/discovery_service_pb2_grpc.py depends on' + + f' grpcio>={GRPC_GENERATED_VERSION}.' + + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}' + + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.' + ) + + +class DiscoveryServiceStub: + """Discovery service used by distributed workers to register their coordinates. + """ + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.Register = channel.unary_unary( + '/tunix.experimental.distributed.runtime.discovery.DiscoveryService/Register', + request_serializer=tunix_dot_experimental_dot_distributed_dot_runtime_dot_discovery_dot_discovery__service__pb2.RegisterRequest.SerializeToString, + response_deserializer=tunix_dot_experimental_dot_distributed_dot_runtime_dot_discovery_dot_discovery__service__pb2.RegisterResponse.FromString, + _registered_method=True) + + +class DiscoveryServiceServicer: + """Discovery service used by distributed workers to register their coordinates. + """ + + def Register(self, request, context): + """Registers a worker node with the discovery service. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + +def add_DiscoveryServiceServicer_to_server(servicer, server): + rpc_method_handlers = { + 'Register': grpc.unary_unary_rpc_method_handler( + servicer.Register, + request_deserializer=tunix_dot_experimental_dot_distributed_dot_runtime_dot_discovery_dot_discovery__service__pb2.RegisterRequest.FromString, + response_serializer=tunix_dot_experimental_dot_distributed_dot_runtime_dot_discovery_dot_discovery__service__pb2.RegisterResponse.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + 'tunix.experimental.distributed.runtime.discovery.DiscoveryService', rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) + server.add_registered_method_handlers('tunix.experimental.distributed.runtime.discovery.DiscoveryService', rpc_method_handlers) + + + # This class is part of an EXPERIMENTAL API. +class DiscoveryService: + """Discovery service used by distributed workers to register their coordinates. + """ + + @staticmethod + def Register(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/tunix.experimental.distributed.runtime.discovery.DiscoveryService/Register', + tunix_dot_experimental_dot_distributed_dot_runtime_dot_discovery_dot_discovery__service__pb2.RegisterRequest.SerializeToString, + tunix_dot_experimental_dot_distributed_dot_runtime_dot_discovery_dot_discovery__service__pb2.RegisterResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) diff --git a/tunix/experimental/examples/math_gsm8k_dist/launch_raiden.sh b/tunix/experimental/examples/math_gsm8k_dist/launch_raiden.sh new file mode 100755 index 000000000..2d862ad94 --- /dev/null +++ b/tunix/experimental/examples/math_gsm8k_dist/launch_raiden.sh @@ -0,0 +1,1183 @@ +#!/bin/bash +# Copyright 2026 Google LLC +# +# Unified Kubernetes JobSet Launcher for Distributed RL Workloads +# (MaxText Trainer + Tunix GRPO Orchestrator + vLLM Rollout + Raiden Weight Sync) +# +# Designed for easy, one-command reproducibility with pre-built container images. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +# Smart resolution of TUNIX repository root and related project directories +if [ -d "${SCRIPT_DIR}/tunix/experimental" ]; then + TUNIX_DIR="${SCRIPT_DIR}" +elif [ -d "${SCRIPT_DIR}/../../../../tunix/experimental" ]; then + TUNIX_DIR="$(cd "${SCRIPT_DIR}/../../../../" && pwd)" +elif [ -d "${SCRIPT_DIR}/tunix" ]; then + TUNIX_DIR="$(cd "${SCRIPT_DIR}/tunix" && pwd)" +else + TUNIX_DIR="${TUNIX_DIR:-$(pwd)}" +fi + +YAML_GENERATOR="${YAML_GENERATOR:-${TUNIX_DIR}/tunix/experimental/distributed/deployment/yaml_generator.py}" +YAMLS_DIR="${YAMLS_DIR:-${TUNIX_DIR}/tunix/experimental/distributed/deployment/yamls}" +MAXTEXT_DIR="${MAXTEXT_DIR:-${TUNIX_DIR}/../maxtext}" +TPU_INFERENCE_DIR="${TPU_INFERENCE_DIR:-${TUNIX_DIR}/../tpu-inference}" + +# ============================================================================== +# Helper: Print Usage +# ============================================================================== +print_usage() { + cat << 'HELP_EOF' +Usage: ./launch_raiden.sh [COMMAND] [OPTIONS] + +Unified script to launch, monitor, triage, and stop distributed RL training +with MaxText, Tunix, vLLM, and Raiden weight sync on GKE TPU clusters. + +COMMANDS: + start (Default) Launch orchestrator, trainer, and rollout JobSets + stop Cleanly delete all JobSets for the current run + restart Stop existing workload and immediately re-launch + status Show running JobSets, pod phases, and node placements + logs [ROLE] [-f] View or follow logs. ROLE: 'trainer' (default), 'rollout', or 'orch' + triage Fetch trainer logs and run automated error diagnosis + dry-run | render Print generated Kubernetes YAML manifests without applying + help, -h, --help Show this help message + +OPTIONS: + --image Container image to run + Default: gcr.io/cloud-tpu-multipod-dev/yixuannwang_google_com-runner:yixuann-raiden-debug-0903-2 + --model, --preset Model preset to use: + - 'qwen3-0.6b' (Default: Qwen3-0.6B, tpuv5:2x2x2 train, tpuv5:2x2x1 roll) + - 'qwen3.5-35b' (Qwen3.5-35B-A3B with scanned ckpt) + - 'qwen3-1.7b' (Qwen3-1.7B standard baseline) + --run-id Custom workload / JobSet prefix + Default: -raiden- + --random-id Append random suffix to run ID for isolated scratch runs + --sync-code Package local tunix/maxtext code to GCS (for rapid live dev) + --no-sync-code Run container image directly as-is (Default: recommended for reproducibility) + --ckpt Custom MaxText checkpoint path (overrides preset default) + --cluster GKE cluster name (default: bodaborg-v5p-nap) + --region GCP region (default: europe-west4) + --zone GCP zone (default: europe-west4-b) + --project GCP project (default: cloud-tpu-shared-capacity) + --cpu-machine CPU machine type for orchestrator (default: n2d-standard-64) + --scratch GCS scratch location + Default: gs://mohitkhatwani_multipods/pathways_scratch/ + --max-steps Max training steps (default: 2) + --max-response-length Max tokens per rollout response (default: 512) + --batch-size Batch size (default: 4) + --rollout-replicas Number of rollout TPU workers (default: 2) + --trainer-slice TPU slice for trainer (default: tpuv5:2x2x2) + --no-cluster-connect Skip automatic gcloud cluster authentication check + --use-ffi Enable Raiden FFI weight synchronization on Pathways TPU workers + --prefuse-moe-weights Prefuse MoE weights (gate + up projection) for rollout TP + --pathways-server-image Pathways server container image (for FFI weight sync) + +EXAMPLES: + # 1. Quick reproduction with verified default image: + ./launch_raiden.sh + + # 2. Launch with a specific given container image: + ./launch_raiden.sh --image gcr.io/cloud-tpu-multipod-dev/my-team-image:tag + + # 3. Check status of running pods: + ./launch_raiden.sh status + + # 4. Stream trainer Python logs: + ./launch_raiden.sh logs trainer -f + + # 5. Automatically diagnose errors in trainer log: + ./launch_raiden.sh triage + + # 6. Stop and delete all JobSets for this run: + ./launch_raiden.sh stop +HELP_EOF +} + +# ============================================================================== +# Model & Workload Presets +# ============================================================================== +load_preset_defaults() { + local preset="$1" + case "$preset" in + qwen3-0.6b|0.6b|0.6B) + PRESET_NAME="qwen3-0.6b" + PRESET_MODEL_TAG="06b" + PRESET_MODEL_NAME="Qwen3-0.6B" + PRESET_MODEL_ID="Qwen/Qwen3-0.6B" + PRESET_MAXTEXT_MODEL_NAME="qwen3-0.6b" + PRESET_TRAINER_BACKEND="maxtext" + PRESET_MAXTEXT_CKPT="gs://maxtext-model-checkpoints/qwen3-0.6b/2025-10-27/scanned/0/items" + PRESET_TRAINER_TPU_SLICE="tpuv5:2x2x2" + PRESET_TRAINER_MESH_FSDP=8 + PRESET_TRAIN_MICRO_BATCH_SIZE=8 + PRESET_BATCH_SIZE=4 + PRESET_NUM_GENERATIONS=2 + PRESET_ROLLOUT_TPU_SLICE="tpuv5:2x2x1" + PRESET_ROLLOUT_TENSOR_PARALLEL_SIZE=2 + PRESET_ROLLOUT_MESH_TP=2 + PRESET_ROLLOUT_REPLICAS=2 + PRESET_SAMPLER="vllm" + PRESET_WEIGHT_SYNC_MODE="raiden" + PRESET_USE_WEIGHT_CONVERTER=1 + PRESET_ROLLOUT_BACKEND="maxtext" + PRESET_VERIFY_WEIGHTS="true" + PRESET_DISABLE_CHECKPOINTING="true" + PRESET_MAX_STEPS=2 + PRESET_DEFAULT_IMAGE="gcr.io/cloud-tpu-multipod-dev/yixuannwang_google_com-runner:yixuann-raiden-debug-0903-2" + PRESET_USE_FFI="false" + PRESET_PREFUSE_MOE_WEIGHTS="false" + PRESET_PATHWAYS_SERVER_IMAGE="us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/gke/datenglin/unsanitized_server:raiden_20260904" + PRESET_PATHWAYS_PROXY_SERVER_IMAGE="us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/gke/datenglin/unsanitized_proxy_server:raiden_20260904" + ;; + qwen3.5-35b|35b|35B) + PRESET_NAME="qwen3.5-35b" + PRESET_MODEL_TAG="35b" + PRESET_MODEL_NAME="Qwen3.5-35B-A3B" + PRESET_MODEL_ID="Qwen/Qwen3.5-35B-A3B" + PRESET_MAXTEXT_MODEL_NAME="qwen3.5-35b-a3b" + PRESET_TRAINER_BACKEND="maxtext" + PRESET_MAXTEXT_CKPT="gs://hengtaoguo-maxtext-logs/checkpoints/qwen3.5-35b-a3b/scanned/2026-06-11-10-27/0/items" + PRESET_TRAINER_TPU_SLICE="tpuv5:2x2x2" + PRESET_TRAINER_MESH_FSDP=4 + PRESET_TRAINER_MESH_TP=2 + PRESET_TRAIN_MICRO_BATCH_SIZE=8 + PRESET_BATCH_SIZE=4 + PRESET_NUM_GENERATIONS=2 + PRESET_ROLLOUT_TPU_SLICE="tpuv5:2x2x1" + PRESET_ROLLOUT_TENSOR_PARALLEL_SIZE=2 + PRESET_ROLLOUT_MESH_TP=2 + PRESET_ROLLOUT_REPLICAS=2 + PRESET_SAMPLER="vllm" + PRESET_WEIGHT_SYNC_MODE="raiden" + PRESET_USE_WEIGHT_CONVERTER=1 + PRESET_ROLLOUT_BACKEND="maxtext" + PRESET_VERIFY_WEIGHTS="true" + PRESET_DISABLE_CHECKPOINTING="true" + PRESET_MAX_STEPS=2 + PRESET_DEFAULT_IMAGE="gcr.io/cloud-tpu-multipod-dev/yixuannwang_google_com-runner:yixuann-raiden-debug-0903-2" + PRESET_USE_FFI="true" + PRESET_PREFUSE_MOE_WEIGHTS="true" + PRESET_PATHWAYS_SERVER_IMAGE="us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/gke/datenglin/unsanitized_server:raiden_20260904" + PRESET_PATHWAYS_PROXY_SERVER_IMAGE="us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/gke/datenglin/unsanitized_proxy_server:raiden_20260904" + ;; + qwen3-1.7b|1.7b|1.7B) + PRESET_NAME="qwen3-1.7b" + PRESET_MODEL_TAG="17b" + PRESET_MODEL_NAME="Qwen3-1.7B" + PRESET_MODEL_ID="Qwen/Qwen3-1.7B" + PRESET_MAXTEXT_MODEL_NAME="qwen3-1.7b" + PRESET_TRAINER_BACKEND="tunix" + PRESET_MAXTEXT_CKPT="" + PRESET_TRAINER_TPU_SLICE="tpuv5:2x2x2" + PRESET_TRAINER_MESH_FSDP=8 + PRESET_TRAIN_MICRO_BATCH_SIZE=1 + PRESET_BATCH_SIZE=4 + PRESET_NUM_GENERATIONS=2 + PRESET_ROLLOUT_TPU_SLICE="tpuv5:2x2x1" + PRESET_ROLLOUT_TENSOR_PARALLEL_SIZE=1 + PRESET_ROLLOUT_MESH_TP=4 + PRESET_ROLLOUT_REPLICAS=1 + PRESET_SAMPLER="vllm" + PRESET_WEIGHT_SYNC_MODE="raiden" + PRESET_USE_WEIGHT_CONVERTER=1 + PRESET_ROLLOUT_BACKEND="maxtext" + PRESET_VERIFY_WEIGHTS="true" + PRESET_DISABLE_CHECKPOINTING="true" + PRESET_MAX_STEPS=1 + PRESET_DEFAULT_IMAGE="gcr.io/cloud-tpu-multipod-dev/yixuannwang_google_com-runner:yixuann-raiden-debug-0903-2" + PRESET_USE_FFI="false" + PRESET_PREFUSE_MOE_WEIGHTS="false" + PRESET_PATHWAYS_SERVER_IMAGE="" + ;; + *) + echo "Error: Unknown preset '$preset'. Available: qwen3-0.6b, qwen3.5-35b, qwen3-1.7b" >&2 + exit 1 + ;; + esac +} + +# ============================================================================== +# Parse Command & CLI Options +# ============================================================================== +COMMAND="" +TARGET_PRESET="${MODEL_PRESET:-${PRESET:-qwen3-0.6b}}" + +USER_IMAGE="" +USER_CKPT="" +USER_RUN_ID="" +USER_MAX_STEPS="" +USER_BATCH_SIZE="" +USER_TRAIN_MICRO_BATCH_SIZE="" +USER_ROLLOUT_REPLICAS="" +USER_TRAINER_TPU_SLICE="" +USER_ROLLOUT_TPU_SLICE="" +USER_SYNC_CODE="" +USER_SCRATCH="" +USER_CLUSTER="" +USER_REGION="" +USER_ZONE="" +USER_PROJECT="" +USER_CPU_MACHINE="" +USER_USE_FFI="" +USER_PATHWAYS_SERVER_IMAGE="" +USER_PATHWAYS_PROXY_SERVER_IMAGE="" +USER_WANDB_API_KEY="" +USER_WANDB_PROJECT="" +USER_WANDB_RUN_NAME="" +USER_WANDB_ENTITY="" +USER_NAMESPACE="" +USER_QUEUE="" +USER_PREFUSE_MOE_WEIGHTS="" +USER_TRAINER_MESH_FSDP="" +USER_TRAINER_MESH_TP="" +USER_USE_WEIGHT_CONVERTER="" + +RANDOMIZE_ID=false +DRY_RUN="${DRY_RUN:-false}" +CONNECT_CLUSTER="${CONNECT_CLUSTER:-true}" +EXTRA_ARGS=() + +while [[ $# -gt 0 ]]; do + case "$1" in + start|stop|restart|status|logs|triage|dry-run|render|start-trainer|start-rollout) + COMMAND="$1" + shift + ;; + help|-h|--help) + print_usage + exit 0 + ;; + --command=*) + COMMAND="${1#*=}" + shift + ;; + --command) + COMMAND="$2" + shift 2 + ;; + --image=*) + USER_IMAGE="${1#*=}" + shift + ;; + --image) + USER_IMAGE="$2" + shift 2 + ;; + --model=*|--preset=*) + TARGET_PRESET="${1#*=}" + shift + ;; + --model|--preset) + TARGET_PRESET="$2" + shift 2 + ;; + --ckpt=*|--maxtext-ckpt=*) + USER_CKPT="${1#*=}" + shift + ;; + --ckpt|--maxtext-ckpt) + USER_CKPT="$2" + shift 2 + ;; + --run-id=*) + USER_RUN_ID="${1#*=}" + shift + ;; + --run-id) + USER_RUN_ID="$2" + shift 2 + ;; + --random-id) + RANDOMIZE_ID=true + shift + ;; + --sync-code) + USER_SYNC_CODE=true + shift + ;; + --no-sync-code) + USER_SYNC_CODE=false + shift + ;; + --use-ffi) + USER_USE_FFI=true + shift + ;; + --no-use-ffi) + USER_USE_FFI=false + shift + ;; + --prefuse-moe-weights=*|--prefuse_moe_weights=*) + USER_PREFUSE_MOE_WEIGHTS="${1#*=}" + shift + ;; + --prefuse-moe-weights|--prefuse_moe_weights) + USER_PREFUSE_MOE_WEIGHTS="$2" + shift 2 + ;; + --no-prefuse-moe-weights) + USER_PREFUSE_MOE_WEIGHTS="false" + shift + ;; + --use-weight-converter|--use_weight_converter) + USER_USE_WEIGHT_CONVERTER=1 + shift + ;; + --no-use-weight-converter|--no_use_weight_converter) + USER_USE_WEIGHT_CONVERTER=0 + shift + ;; + --max-response-length=*|--max_response_length=*) + MAX_RESPONSE_LENGTH="${1#*=}" + shift + ;; + --max-response-length|--max_response_length) + MAX_RESPONSE_LENGTH="$2" + shift 2 + ;; + --pathways-server-image=*) + USER_PATHWAYS_SERVER_IMAGE="${1#*=}" + shift + ;; + --pathways-server-image) + USER_PATHWAYS_SERVER_IMAGE="$2" + shift 2 + ;; + --pathways-proxy-server-image=*|--pathways-proxy-image=*) + USER_PATHWAYS_PROXY_SERVER_IMAGE="${1#*=}" + shift + ;; + --pathways-proxy-server-image|--pathways-proxy-image) + USER_PATHWAYS_PROXY_SERVER_IMAGE="$2" + shift 2 + ;; + --cluster=*) + USER_CLUSTER="${1#*=}" + shift + ;; + --cluster) + USER_CLUSTER="$2" + shift 2 + ;; + --region=*) + USER_REGION="${1#*=}" + shift + ;; + --region) + USER_REGION="$2" + shift 2 + ;; + --zone=*) + USER_ZONE="${1#*=}" + shift + ;; + --zone) + USER_ZONE="$2" + shift 2 + ;; + --project=*) + USER_PROJECT="${1#*=}" + shift + ;; + --project) + USER_PROJECT="$2" + shift 2 + ;; + --cpu-machine=*) + USER_CPU_MACHINE="${1#*=}" + shift + ;; + --cpu-machine) + USER_CPU_MACHINE="$2" + shift 2 + ;; + --scratch=*) + USER_SCRATCH="${1#*=}" + shift + ;; + --scratch) + USER_SCRATCH="$2" + shift 2 + ;; + --max-steps=*) + USER_MAX_STEPS="${1#*=}" + shift + ;; + --max-steps) + USER_MAX_STEPS="$2" + shift 2 + ;; + --batch-size=*) + USER_BATCH_SIZE="${1#*=}" + shift + ;; + --batch-size) + USER_BATCH_SIZE="$2" + shift 2 + ;; + --train-micro-batch-size=*) + USER_TRAIN_MICRO_BATCH_SIZE="${1#*=}" + shift + ;; + --train-micro-batch-size) + USER_TRAIN_MICRO_BATCH_SIZE="$2" + shift 2 + ;; + --rollout-replicas=*) + USER_ROLLOUT_REPLICAS="${1#*=}" + shift + ;; + --rollout-replicas) + USER_ROLLOUT_REPLICAS="$2" + shift 2 + ;; + --trainer-slice=*) + USER_TRAINER_TPU_SLICE="${1#*=}" + shift + ;; + --trainer-slice) + USER_TRAINER_TPU_SLICE="$2" + shift 2 + ;; + --rollout-slice=*) + USER_ROLLOUT_TPU_SLICE="${1#*=}" + shift + ;; + --rollout-slice) + USER_ROLLOUT_TPU_SLICE="$2" + shift 2 + ;; + --dry-run) + DRY_RUN=true + shift + ;; + --namespace=*) + USER_NAMESPACE="${1#*=}" + shift + ;; + --namespace) + USER_NAMESPACE="$2" + shift 2 + ;; + --queue=*|--kueue-queue=*) + USER_QUEUE="${1#*=}" + shift + ;; + --queue|--kueue-queue) + USER_QUEUE="$2" + shift 2 + ;; + --no-cluster-connect) + CONNECT_CLUSTER=false + shift + ;; + --verify-weights) + USER_VERIFY_WEIGHTS="true" + shift + ;; + --no-verify-weights) + USER_VERIFY_WEIGHTS="false" + shift + ;; + --trainer-mesh-fsdp=*) + USER_TRAINER_MESH_FSDP="${1#*=}" + shift + ;; + --trainer-mesh-fsdp) + USER_TRAINER_MESH_FSDP="$2" + shift 2 + ;; + --trainer-mesh-tp=*) + USER_TRAINER_MESH_TP="${1#*=}" + shift + ;; + --trainer-mesh-tp) + USER_TRAINER_MESH_TP="$2" + shift 2 + ;; + --sync-tpu-inference) + USER_SYNC_TPU_INFERENCE="true" + shift + ;; + --no-sync-tpu-inference) + USER_SYNC_TPU_INFERENCE="false" + shift + ;; + --wandb-api-key=*) + USER_WANDB_API_KEY="${1#*=}" + shift + ;; + --wandb-api-key) + USER_WANDB_API_KEY="$2" + shift 2 + ;; + --wandb-project=*) + USER_WANDB_PROJECT="${1#*=}" + shift + ;; + --wandb-project) + USER_WANDB_PROJECT="$2" + shift 2 + ;; + --wandb-run-name=*) + USER_WANDB_RUN_NAME="${1#*=}" + shift + ;; + --wandb-run-name) + USER_WANDB_RUN_NAME="$2" + shift 2 + ;; + --wandb-entity=*) + USER_WANDB_ENTITY="${1#*=}" + shift + ;; + --wandb-entity) + USER_WANDB_ENTITY="$2" + shift 2 + ;; + *) + EXTRA_ARGS+=("$1") + shift + ;; + esac +done + +# Load selected preset values +load_preset_defaults "${TARGET_PRESET}" + +# Resolve all final configuration values: CLI flag > Env Var > Preset Default +MODEL_NAME="${MODEL_NAME:-${PRESET_MODEL_NAME}}" +MODEL_ID="${MODEL_ID:-${PRESET_MODEL_ID}}" +MODEL_TAG="${MODEL_TAG:-${PRESET_MODEL_TAG}}" +MAXTEXT_MODEL_NAME="${MAXTEXT_MODEL_NAME:-${PRESET_MAXTEXT_MODEL_NAME}}" +TRAINER_BACKEND="${TRAINER_BACKEND:-${PRESET_TRAINER_BACKEND}}" +TOKENIZER_PATH="${TOKENIZER_PATH:-${MODEL_ID}}" + +MAXTEXT_CKPT="${USER_CKPT:-${MAXTEXT_CKPT:-${PRESET_MAXTEXT_CKPT}}}" +TUNIX_IMAGE="${USER_IMAGE:-${TUNIX_IMAGE:-${PRESET_DEFAULT_IMAGE}}}" +MAX_STEPS="${USER_MAX_STEPS:-${MAX_STEPS:-${PRESET_MAX_STEPS}}}" +BATCH_SIZE="${USER_BATCH_SIZE:-${BATCH_SIZE:-${PRESET_BATCH_SIZE}}}" +NUM_GENERATIONS="${NUM_GENERATIONS:-${PRESET_NUM_GENERATIONS}}" +TRAIN_MICRO_BATCH_SIZE="${USER_TRAIN_MICRO_BATCH_SIZE:-${TRAIN_MICRO_BATCH_SIZE:-${PRESET_TRAIN_MICRO_BATCH_SIZE}}}" +ROLLOUT_REPLICAS="${USER_ROLLOUT_REPLICAS:-${ROLLOUT_REPLICAS:-${PRESET_ROLLOUT_REPLICAS}}}" +TRAINER_TPU_SLICE="${USER_TRAINER_TPU_SLICE:-${TRAINER_TPU_SLICE:-${PRESET_TRAINER_TPU_SLICE}}}" +ROLLOUT_TPU_SLICE="${USER_ROLLOUT_TPU_SLICE:-${ROLLOUT_TPU_SLICE:-${PRESET_ROLLOUT_TPU_SLICE}}}" +TRAINER_MESH_FSDP="${USER_TRAINER_MESH_FSDP:-${TRAINER_MESH_FSDP:-${PRESET_TRAINER_MESH_FSDP}}}" +ROLLOUT_MESH_TP="${ROLLOUT_MESH_TP:-${PRESET_ROLLOUT_MESH_TP}}" +SAMPLER="${SAMPLER:-${PRESET_SAMPLER}}" +WEIGHT_SYNC_MODE="${WEIGHT_SYNC_MODE:-${PRESET_WEIGHT_SYNC_MODE}}" +USE_WEIGHT_CONVERTER="${USER_USE_WEIGHT_CONVERTER:-${USE_WEIGHT_CONVERTER:-${PRESET_USE_WEIGHT_CONVERTER}}}" +ROLLOUT_BACKEND="${ROLLOUT_BACKEND:-${PRESET_ROLLOUT_BACKEND}}" +VERIFY_WEIGHTS="${USER_VERIFY_WEIGHTS:-${VERIFY_WEIGHTS:-${PRESET_VERIFY_WEIGHTS}}}" +DISABLE_CHECKPOINTING="${DISABLE_CHECKPOINTING:-${PRESET_DISABLE_CHECKPOINTING}}" +USE_FFI="${USER_USE_FFI:-${USE_FFI:-${PRESET_USE_FFI:-false}}}" +PREFUSE_MOE_WEIGHTS="${USER_PREFUSE_MOE_WEIGHTS:-${PREFUSE_MOE_WEIGHTS:-${PRESET_PREFUSE_MOE_WEIGHTS:-false}}}" +DEFAULT_FFI_SERVER_IMAGE="us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/gke/datenglin/unsanitized_server:raiden_20260904" +DEFAULT_FFI_PROXY_IMAGE="us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/gke/datenglin/unsanitized_proxy_server:raiden_20260904" +if [[ "${USE_FFI}" == "true" ]]; then + PATHWAYS_SERVER_IMAGE="${USER_PATHWAYS_SERVER_IMAGE:-${PATHWAYS_SERVER_IMAGE:-${PRESET_PATHWAYS_SERVER_IMAGE:-${DEFAULT_FFI_SERVER_IMAGE}}}}" + PATHWAYS_PROXY_SERVER_IMAGE="${USER_PATHWAYS_PROXY_SERVER_IMAGE:-${PATHWAYS_PROXY_SERVER_IMAGE:-${PRESET_PATHWAYS_PROXY_SERVER_IMAGE:-${DEFAULT_FFI_PROXY_IMAGE}}}}" +else + PATHWAYS_SERVER_IMAGE="${USER_PATHWAYS_SERVER_IMAGE:-${PATHWAYS_SERVER_IMAGE:-${PRESET_PATHWAYS_SERVER_IMAGE:-}}}" + PATHWAYS_PROXY_SERVER_IMAGE="${USER_PATHWAYS_PROXY_SERVER_IMAGE:-${PATHWAYS_PROXY_SERVER_IMAGE:-${PRESET_PATHWAYS_PROXY_SERVER_IMAGE:-}}}" +fi + +SYNC_CODE="${USER_SYNC_CODE:-${SYNC_CODE:-true}}" # Default: true to patch local tunix/maxtext changes +SYNC_TPU_INFERENCE="${USER_SYNC_TPU_INFERENCE:-${SYNC_TPU_INFERENCE:-true}}" +PROJECT="${USER_PROJECT:-${PROJECT:-cloud-tpu-shared-capacity}}" +REGION="${USER_REGION:-${REGION:-europe-west4}}" +ZONE="${USER_ZONE:-${ZONE:-europe-west4-b}}" +CLUSTER="${USER_CLUSTER:-${CLUSTER:-bodaborg-v5p-nap}}" +CPU_MACHINE="${USER_CPU_MACHINE:-${CPU_MACHINE:-n2d-standard-64}}" + +CURRENT_SYSTEM_USER="${USER:-$(whoami 2>/dev/null || echo "user")}" +if [[ -n "${USER_RUN_ID}" ]]; then + RUN_ID="${USER_RUN_ID}" +elif [[ "${RANDOMIZE_ID}" == "true" ]]; then + RUN_ID="${CURRENT_SYSTEM_USER:0:8}-r$((RANDOM % 90000 + 10000))" +else + # JobSet coordinator label has format "-proc-0-0.". + # Kubernetes label values are strictly limited to 63 characters. + # 2 * len(TRAINER_ID) + 10 <= 63 => len(TRAINER_ID) <= 26. + # TRAINER_ID is "${RUN_ID}-train", so len(RUN_ID) must be <= 20. + RUN_ID="${CURRENT_SYSTEM_USER:0:8}-rd-${MODEL_TAG}" +fi + +export USER="${RUN_ID}" +GCS_SCRATCH_LOCATION="${USER_SCRATCH:-${GCS_SCRATCH_LOCATION:-gs://mohitkhatwani_multipods/pathways_scratch/${RUN_ID}}}" +GCS_SYNC_TAR="${GCS_SCRATCH_LOCATION}/code_sync/${USER}.tar.gz" + +ORCHESTRATOR_ID="${RUN_ID}-orch" +TRAINER_ID="${RUN_ID}-train" +ROLLOUT_ID="${RUN_ID}-roll" + +if [[ ${#TRAINER_ID} -gt 26 ]]; then + echo "Error: TRAINER_ID '${TRAINER_ID}' is too long (${#TRAINER_ID} chars). Maximum allowed is 26 characters due to Kubernetes 63-char label value limit on JobSet coordinator label (${TRAINER_ID}-proc-0-0.${TRAINER_ID}). Please specify a shorter --run-id." >&2 + exit 1 +fi + +NAMESPACE="${USER_NAMESPACE:-${K8S_NAMESPACE:-default}}" +QUEUE_NAME="${USER_QUEUE:-${KUEUE_QUEUE:-}}" + +MODEL_DIR="${MODEL_DIR:-}" +MAX_PROMPT_LENGTH="${MAX_PROMPT_LENGTH:-512}" +MAX_RESPONSE_LENGTH="${MAX_RESPONSE_LENGTH:-512}" +MINI_BATCH_SIZE="${MINI_BATCH_SIZE:-$((BATCH_SIZE * NUM_GENERATIONS))}" +EVAL_EVERY_N_STEPS="${EVAL_EVERY_N_STEPS:-1000000}" +LORA_RANK="${LORA_RANK:-16}" +LORA_ALPHA="${LORA_ALPHA:-16.0}" +DEBUG="${DEBUG:-0}" + +MAXTEXT_OUTPUT_DIR="${MAXTEXT_OUTPUT_DIR:-/tmp/artifacts/math_gsm8k_dist/maxtext}" +TRAINER_MESH_TP="${USER_TRAINER_MESH_TP:-${TRAINER_MESH_TP:-${PRESET_TRAINER_MESH_TP:-1}}}" +TRAINER_MESH_EXPERT="${TRAINER_MESH_EXPERT:-1}" +TRAINER_PADDED_MOE_MLP_DIM="${TRAINER_PADDED_MOE_MLP_DIM:-}" +ROLLOUT_USE_BATCHED_RPA="${ROLLOUT_USE_BATCHED_RPA:-}" +ROLLOUT_MAXTEXT_ATTENTION="${ROLLOUT_MAXTEXT_ATTENTION:-}" +TRAINER_JOBSET_YAML="${TRAINER_JOBSET_YAML:-jobset.pathways.yaml}" + +WANDB_API_KEY="${USER_WANDB_API_KEY:-${WANDB_API_KEY:-}}" +WANDB_PROJECT="${USER_WANDB_PROJECT:-${WANDB_PROJECT:-trellis-gsm8k}}" +WANDB_RUN_NAME="${USER_WANDB_RUN_NAME:-${WANDB_RUN_NAME:-}}" +WANDB_ENTITY="${USER_WANDB_ENTITY:-${WANDB_ENTITY:-google-trellis}}" + +ORCHESTRATOR_PORT="${ORCHESTRATOR_PORT:-20000}" +ROLLOUT_PORT="${ROLLOUT_PORT:-20001}" +TRAINER_PORT="${TRAINER_PORT:-20002}" + +COMMAND="${COMMAND:-start}" +if [[ "$COMMAND" == "dry-run" || "$COMMAND" == "render" ]]; then + COMMAND="start" + DRY_RUN=true +fi + +# ============================================================================== +# Cluster Connection +# ============================================================================== +connect_cluster() { + if [[ "${CONNECT_CLUSTER}" != "true" || "${DRY_RUN}" == "true" ]]; then + return 0 + fi + + local target_context="gke_${PROJECT}_${REGION}_${CLUSTER}" + local current_context + current_context=$(kubectl config current-context 2>/dev/null || true) + + if [[ "$current_context" == "$target_context" ]]; then + return 0 + fi + + echo "=================================================================" + echo "Connecting to GKE cluster: ${CLUSTER} in ${REGION} (${PROJECT})..." + echo "=================================================================" + gcloud container clusters get-credentials "${CLUSTER}" \ + --region "${REGION}" \ + --project "${PROJECT}" \ + --dns-endpoint + kubectl config use-context "${target_context}" >/dev/null 2>&1 || true +} + +# ============================================================================== +# Code Sync Logic (Optional for given-image runs) +# ============================================================================== +sync_code_to_gcs() { + if [[ "${SYNC_CODE}" != "true" ]]; then + echo "ℹī¸ Code sync disabled (SYNC_CODE=false). Using container image as-is." + return 0 + fi + + if [ ! -d "${TUNIX_DIR}" ] || [ ! -d "${MAXTEXT_DIR}" ]; then + echo "Error: Local tunix/ or maxtext/ directories not found for code sync." >&2 + echo "Expected at: ${TUNIX_DIR} and ${MAXTEXT_DIR}" >&2 + exit 1 + fi + + local pkg_msg="tunix/, maxtext/" + local extra_dirs=() + if [ -d "${MAXTEXT_DIR}" ]; then + local maxtext_parent + maxtext_parent="$(cd "${MAXTEXT_DIR}/.." && pwd)" + extra_dirs+=("-C" "${maxtext_parent}" "maxtext") + fi + if [[ "${SYNC_TPU_INFERENCE:-true}" == "true" ]] && [ -d "${TPU_INFERENCE_DIR}" ]; then + local tpu_inf_parent + tpu_inf_parent="$(cd "${TPU_INFERENCE_DIR}/.." && pwd)" + extra_dirs+=("-C" "${tpu_inf_parent}" "tpu-inference") + pkg_msg="tunix/, maxtext/, tpu-inference/" + fi + + echo "=================================================================" + echo "đŸ“Ļ Packaging local changes from ${pkg_msg}..." + echo "=================================================================" + local tar_file="/tmp/code_sync_${USER}.tar.gz" + rm -f "${tar_file}" + + tar --exclude=".git" \ + --exclude=".venv" \ + --exclude="venv*" \ + --exclude="myenv" \ + --exclude="__pycache__" \ + --exclude=".pytest_cache" \ + --exclude="docs" \ + --exclude="benchmarks" \ + --exclude="tests" \ + --exclude="artifacts" \ + --exclude="checkpoints" \ + -czf "${tar_file}" \ + -C "${TUNIX_DIR}" . \ + "${extra_dirs[@]}" + + local tar_size + tar_size=$(ls -lh "${tar_file}" | awk '{print $5}') + echo "đŸ“Ļ Archive created (${tar_size}). Uploading to ${GCS_SYNC_TAR}..." + + if [[ "${DRY_RUN}" == "true" ]]; then + echo "[DRY RUN] Would upload ${tar_file} to ${GCS_SYNC_TAR}" + else + if command -v gcloud &>/dev/null; then + gcloud storage cp "${tar_file}" "${GCS_SYNC_TAR}" --quiet || gsutil cp "${tar_file}" "${GCS_SYNC_TAR}" + else + gsutil cp "${tar_file}" "${GCS_SYNC_TAR}" + fi + echo "✅ Local code synced to GCS: ${GCS_SYNC_TAR}" + fi + + rm -f "${tar_file}" + echo "=================================================================" +} + +get_sync_prefix() { + if [[ "${SYNC_CODE}" == "true" ]]; then + local python_path="/app/maxtext/src:/app/tpu-inference:/app:\${PYTHONPATH:-}" + echo "echo '==> Syncing code from ${GCS_SYNC_TAR}...'; (gcloud storage cp ${GCS_SYNC_TAR} /tmp/code_sync.tar.gz 2>/dev/null || gsutil cp ${GCS_SYNC_TAR} /tmp/code_sync.tar.gz 2>/dev/null || python3 -c \"import json, urllib.request, urllib.parse; token = json.loads(urllib.request.urlopen(urllib.request.Request('http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/token', headers={'Metadata-Flavor': 'Google'})).read())['access_token']; b, p = '${GCS_SYNC_TAR}'.replace('gs://', '').split('/', 1); req = urllib.request.Request(f'https://storage.googleapis.com/storage/v1/b/{b}/o/{urllib.parse.quote(p, safe=\\\"\\\")}?alt=media', headers={'Authorization': f'Bearer {token}'}); open('/tmp/code_sync.tar.gz', 'wb').write(urllib.request.urlopen(req).read())\" 2>/dev/null || python3 -c \"import gcsfs; gcsfs.GCSFileSystem().get('${GCS_SYNC_TAR}', '/tmp/code_sync.tar.gz')\") && tar -xzf /tmp/code_sync.tar.gz -C /app && rm -f /tmp/code_sync.tar.gz && (pip install --force-reinstall --no-deps /app/raiden_wheels/*.whl 2>/dev/null || true) && (cp -rf /app/tunix /opt/venv/lib/python3.12/site-packages/ 2>/dev/null || true) && (cp -rf /app/maxtext/src/maxtext /opt/venv/lib/python3.12/site-packages/ 2>/dev/null || true) && ([ -d /app/tpu-inference/tpu_inference ] && cp -rf /app/tpu-inference/tpu_inference /opt/venv/lib/python3.12/site-packages/ 2>/dev/null || true) && export PYTHONPATH=\"${python_path}\" && echo '==> Code sync and Raiden wheel applied to /app and site-packages.';" + else + echo "" + fi +} + +apply_or_print() { + if [[ "${DRY_RUN}" == "true" ]]; then + echo "---" + cat + else + kubectl apply -n "${NAMESPACE}" -f - + fi +} + +# ============================================================================== +# JobSet Lifecycle Operations +# ============================================================================== +stop_workload() { + local jobsets=("${ORCHESTRATOR_ID}" "${TRAINER_ID}" "${ROLLOUT_ID}") + for (( i=0; i<8; i++ )); do + jobsets+=("${ROLLOUT_ID}-${i}") + done + + if [[ "${DRY_RUN}" == "true" ]]; then + echo "[DRY RUN] Would delete jobsets: ${jobsets[*]}" + return 0 + fi + + echo "Stopping workload (${RUN_ID}) in namespace ${NAMESPACE}..." + for js in "${jobsets[@]}"; do + if kubectl get jobset "$js" -n "${NAMESPACE}" &>/dev/null; then + echo " Deleting JobSet: $js in namespace ${NAMESPACE}..." + kubectl delete jobset "$js" -n "${NAMESPACE}" --ignore-not-found=true --wait=true 2>/dev/null || true + kubectl delete jobs -n "${NAMESPACE}" -l "jobset.sigs.k8s.io/jobset-name=$js" --ignore-not-found=true 2>/dev/null || true + fi + done + echo "✅ Teardown command sent for ${RUN_ID}." +} + +start_orchestrator() { + local sync_prefix + sync_prefix=$(get_sync_prefix) + + echo "Rendering & Starting orchestrator (${ORCHESTRATOR_ID}) in namespace ${NAMESPACE}..." + python3 "${YAML_GENERATOR}" \ + "${YAMLS_DIR}/jobset.cpu.yaml" \ + --jobset_name="${ORCHESTRATOR_ID}" \ + --namespace="${NAMESPACE}" \ + ${QUEUE_NAME:+--queue_name="${QUEUE_NAME}"} \ + --cpu_machine="${CPU_MACHINE}" \ + --worker_container_image="${TUNIX_IMAGE}" \ + --worker_container_port="${ORCHESTRATOR_PORT}" \ + --worker_startup_command=" \ + ${sync_prefix} \ + ${WANDB_API_KEY:+WANDB_API_KEY=\"${WANDB_API_KEY}\"} \ + ${WANDB_PROJECT:+WANDB_PROJECT=\"${WANDB_PROJECT}\"} \ + ${WANDB_RUN_NAME:+WANDB_RUN_NAME=\"${WANDB_RUN_NAME}\"} \ + ${WANDB_ENTITY:+WANDB_ENTITY=\"${WANDB_ENTITY}\"} \ + PYTHONUNBUFFERED=1 DISABLE_CHECKPOINTING=${DISABLE_CHECKPOINTING} python -m tunix.experimental.distributed.runtime.main \ + --discovery_id=${ORCHESTRATOR_ID} \ + --discovery_port=${ORCHESTRATOR_PORT} \ + --process_main=tunix.experimental.examples.math_gsm8k_dist.run_gsm8k_dist_grpo.main \ + --model_id=${MODEL_ID} \ + --tokenizer_path=${TOKENIZER_PATH} \ + --batch_size=${BATCH_SIZE} \ + --num_generations=${NUM_GENERATIONS} \ + --max_steps=${MAX_STEPS} \ + --max_prompt_length=${MAX_PROMPT_LENGTH} \ + --max_response_length=${MAX_RESPONSE_LENGTH} \ + --train_micro_batch_size=${TRAIN_MICRO_BATCH_SIZE} \ + --rollout_replicas=${ROLLOUT_REPLICAS} \ + --wandb_project=\"${WANDB_PROJECT}\" \ + --wandb_run_name=\"${WANDB_RUN_NAME}\" \ + --weight_sync_mode=${WEIGHT_SYNC_MODE} \ + --stop_workers_on_exit \ + ${DEBUG:+--debug} \ + " \ + | apply_or_print +} + +start_trainer() { + local maxtext_args="" + if [[ "${TRAINER_BACKEND}" == "maxtext" ]]; then + maxtext_args=" \ + --maxtext_model_name=${MAXTEXT_MODEL_NAME} \ + ${TRAINER_PADDED_MOE_MLP_DIM:+--maxtext_padded_moe_mlp_dim=${TRAINER_PADDED_MOE_MLP_DIM}} \ + --maxtext_ckpt_path=${MAXTEXT_CKPT} \ + --maxtext_output_directory=${MAXTEXT_OUTPUT_DIR} \ + --mesh_tp=${TRAINER_MESH_TP} \ + --mesh_expert=${TRAINER_MESH_EXPERT} \ + --prefuse_moe_weights=${PREFUSE_MOE_WEIGHTS} \ + " + fi + + local sync_prefix + sync_prefix=$(get_sync_prefix) + + local ffi_env="" + if [[ "${USE_FFI}" == "true" ]]; then + ffi_env="USE_RAIDEN_FFI=true RAIDEN_USE_FFI=1 RAIDEN_DEVICES_PER_HOST=${RAIDEN_DEVICES_PER_HOST:-4}" + else + ffi_env="USE_RAIDEN_FFI=false RAIDEN_USE_FFI=0" + fi + + local pw_server_arg=() + if [[ -n "${PATHWAYS_SERVER_IMAGE}" ]]; then + pw_server_arg+=("--pathways_server_image=${PATHWAYS_SERVER_IMAGE}") + fi + if [[ -n "${PATHWAYS_PROXY_SERVER_IMAGE}" ]]; then + pw_server_arg+=("--pathways_proxy_server_image=${PATHWAYS_PROXY_SERVER_IMAGE}") + fi + + echo "Rendering & Starting trainer (${TRAINER_ID}) in namespace ${NAMESPACE}..." + python3 "${YAML_GENERATOR}" \ + "${YAMLS_DIR}/${TRAINER_JOBSET_YAML}" \ + --jobset_name="${TRAINER_ID}" \ + --namespace="${NAMESPACE}" \ + ${QUEUE_NAME:+--queue_name="${QUEUE_NAME}"} \ + --tpu_slice="${TRAINER_TPU_SLICE}" \ + --cpu_machine="${CPU_MACHINE}" \ + --pathways_gcs_scratch_location="${GCS_SCRATCH_LOCATION}" \ + --worker_container_image="${TUNIX_IMAGE}" \ + --worker_container_port="${TRAINER_PORT}" \ + "${pw_server_arg[@]}" \ + --worker_startup_command=" \ + ${sync_prefix} \ + PYTHONUNBUFFERED=1 ${ffi_env} PREFUSE_MOE_WEIGHTS=${PREFUSE_MOE_WEIGHTS} ROLLOUT_TENSOR_PARALLEL_SIZE=${ROLLOUT_MESH_TP} DISABLE_CHECKPOINTING=${DISABLE_CHECKPOINTING} VERIFY_WEIGHTS=${VERIFY_WEIGHTS} USE_WEIGHT_CONVERTER=${USE_WEIGHT_CONVERTER} ROLLOUT_BACKEND=${ROLLOUT_BACKEND} python -m tunix.experimental.distributed.runtime.main \ + --discovery_addrs=${ORCHESTRATOR_ID}:${ORCHESTRATOR_PORT} \ + --process_executor=tunix.experimental.distributed.runtime.executor.K8sExecutor \ + --process_main=tunix.experimental.examples.math_gsm8k_dist.run_trainer_node.main \ + --worker_id=${TRAINER_ID} \ + --port=${TRAINER_PORT} \ + --mesh_fsdp=${TRAINER_MESH_FSDP} \ + --trainer_backend=${TRAINER_BACKEND} \ + --model_name=${MODEL_NAME} \ + --model_id=${MODEL_ID} \ + --model_dir=${MODEL_DIR} \ + --tokenizer_path=${TOKENIZER_PATH} \ + --max_prompt_length=${MAX_PROMPT_LENGTH} \ + --max_response_length=${MAX_RESPONSE_LENGTH} \ + --mini_batch_size=${MINI_BATCH_SIZE} \ + --train_micro_batch_size=${TRAIN_MICRO_BATCH_SIZE} \ + --eval_every_n_steps=${EVAL_EVERY_N_STEPS} \ + --lora_rank=${LORA_RANK} \ + --lora_alpha=${LORA_ALPHA} \ + ${ROLLOUT_MESH_TP:+--rollout_mesh_tp=${ROLLOUT_MESH_TP}} \ + ${maxtext_args} \ + " \ + | apply_or_print +} + +start_rollout_instance() { + local target_id="$1" + local maxtext_args="" + if [[ "${TRAINER_BACKEND}" == "maxtext" ]]; then + maxtext_args=" \ + --maxtext_model_name=${MAXTEXT_MODEL_NAME} \ + ${ROLLOUT_MAXTEXT_ATTENTION:+--maxtext_attention=${ROLLOUT_MAXTEXT_ATTENTION}} \ + --prefuse_moe_weights=${PREFUSE_MOE_WEIGHTS} \ + " + fi + local vllm_args="" + if [[ "$SAMPLER" == "vllm" ]]; then + vllm_args="\ + --sampler=vllm \ + --sampler_mesh_tp=${ROLLOUT_MESH_TP} \ + --mesh_tp=${ROLLOUT_MESH_TP} \ + " + fi + + local sync_prefix + sync_prefix=$(get_sync_prefix) + + local rollout_ffi_env="USE_RAIDEN_FFI=false RAIDEN_USE_FFI=0" + + echo "Rendering & Starting rollout (${target_id}) in namespace ${NAMESPACE}..." + python3 "${YAML_GENERATOR}" \ + "${YAMLS_DIR}/jobset.tpu.yaml" \ + --jobset_name="${target_id}" \ + --namespace="${NAMESPACE}" \ + ${QUEUE_NAME:+--queue_name="${QUEUE_NAME}"} \ + --tpu_slice="${ROLLOUT_TPU_SLICE}" \ + --pathways_gcs_scratch_location="${GCS_SCRATCH_LOCATION}" \ + --worker_container_image="${TUNIX_IMAGE}" \ + --worker_container_port="${ROLLOUT_PORT}" \ + --worker_startup_command=" \ + ${sync_prefix} \ + PYTHONUNBUFFERED=1 ${rollout_ffi_env} PREFUSE_MOE_WEIGHTS=${PREFUSE_MOE_WEIGHTS} ROLLOUT_TENSOR_PARALLEL_SIZE=${ROLLOUT_MESH_TP} SKIP_JAX_PRECOMPILE=1 VERIFY_WEIGHTS=${VERIFY_WEIGHTS} ${ROLLOUT_USE_BATCHED_RPA:+USE_BATCHED_RPA_KERNEL=1} python -m tunix.experimental.distributed.runtime.main \ + --discovery_addrs=${ORCHESTRATOR_ID}:${ORCHESTRATOR_PORT} \ + --process_executor=tunix.experimental.distributed.runtime.executor.K8sExecutor \ + --process_main=tunix.experimental.examples.math_gsm8k_dist.run_rollout_node.main \ + --worker_id=${target_id} \ + --port=${ROLLOUT_PORT} \ + --model_name=${MODEL_NAME} \ + --model_id=${MODEL_ID} \ + --model_dir=${MODEL_DIR} \ + --tokenizer_path=${TOKENIZER_PATH} \ + --max_prompt_length=${MAX_PROMPT_LENGTH} \ + --max_response_length=${MAX_RESPONSE_LENGTH} \ + --lora_rank=${LORA_RANK} \ + --lora_alpha=${LORA_ALPHA} \ + --weight_sync_mode=${WEIGHT_SYNC_MODE} \ + ${maxtext_args} \ + ${vllm_args} \ + ${DEBUG:+--debug} \ + " \ + | apply_or_print +} + +start_all_rollouts() { + if [[ "${ROLLOUT_REPLICAS}" -gt 1 ]]; then + for (( i=0; i/dev/null || true) + local js_output + js_output=$(echo "$js_all" | grep "${RUN_ID}" || true) + if [[ -n "$js_output" ]]; then + echo "$js_all" | { head -n 1 || true; } + echo "$js_output" + else + echo "No active JobSets found for ${RUN_ID} in namespace ${NAMESPACE}." + fi + echo "" + echo "================================================================================" + echo "Pods for Run: ${RUN_ID} in namespace ${NAMESPACE}" + echo "================================================================================" + local pod_all + pod_all=$(kubectl get pods -n "${NAMESPACE}" -o wide 2>/dev/null || true) + local pod_output + pod_output=$(echo "$pod_all" | grep "${RUN_ID}" || true) + if [[ -n "$pod_output" ]]; then + echo "$pod_all" | { head -n 1 || true; } + echo "$pod_output" + else + echo "No active Pods found for ${RUN_ID} in namespace ${NAMESPACE}." + fi +} + +show_logs() { + connect_cluster + local role="${1:-trainer}" + shift || true + local log_args=("$@") + + local target_jobset="" + local container_flag="-c main" + + case "$role" in + trainer|train) + target_jobset="${TRAINER_ID}" + ;; + rollout|roll) + if [[ "${ROLLOUT_REPLICAS}" -gt 1 ]]; then + target_jobset="${ROLLOUT_ID}-0" + else + target_jobset="${ROLLOUT_ID}" + fi + ;; + orchestrator|orch) + target_jobset="${ORCHESTRATOR_ID}" + ;; + *) + echo "Unknown role '$role'. Valid options: trainer, rollout, orch" >&2 + exit 1 + ;; + esac + + local pod_name + pod_name=$(kubectl get pods -n "${NAMESPACE}" -l "jobset.sigs.k8s.io/jobset-name=${target_jobset},jobset.sigs.k8s.io/replicatedjob-name=proc" -o jsonpath='{.items[0].metadata.name}' 2>/dev/null || true) + if [[ -z "$pod_name" ]]; then + pod_name=$(kubectl get pods -n "${NAMESPACE}" -l "jobset.sigs.k8s.io/jobset-name=${target_jobset}" -o jsonpath='{.items[0].metadata.name}' 2>/dev/null || true) + fi + + if [[ -z "$pod_name" ]]; then + echo "Error: No pod found for JobSet ${target_jobset} in namespace ${NAMESPACE}." >&2 + echo "Check current status with: $0 status --namespace=${NAMESPACE}" >&2 + exit 1 + fi + + echo "Fetching logs from pod ${pod_name} (${container_flag}) in namespace ${NAMESPACE}..." + kubectl logs -n "${NAMESPACE}" "${pod_name}" ${container_flag} "${log_args[@]}" +} + +triage_workload() { + connect_cluster + local pod_name + pod_name=$(kubectl get pods -n "${NAMESPACE}" -l "jobset.sigs.k8s.io/jobset-name=${TRAINER_ID},jobset.sigs.k8s.io/replicatedjob-name=proc" -o jsonpath='{.items[0].metadata.name}' 2>/dev/null || true) + + if [[ -z "$pod_name" ]]; then + echo "Error: Trainer proc pod not found for ${TRAINER_ID} in namespace ${NAMESPACE}." >&2 + echo "Workload might not be running. Run '$0 status --namespace=${NAMESPACE}' to inspect." >&2 + exit 1 + fi + + local log_file="/tmp/triage-${RUN_ID}-trainer.log" + echo "Downloading recent trainer logs from ${pod_name} to ${log_file} in namespace ${NAMESPACE}..." + kubectl logs -n "${NAMESPACE}" "${pod_name}" -c main --tail=1000 > "${log_file}" 2>&1 || true + + echo "================================================================================" + echo "TRIAGE REPORT FOR: ${RUN_ID}" + echo "================================================================================" + + if grep -E "(WeightConverterError|ConversionPlanError|ShapeMismatchError|Push range out of bounds)" "${log_file}"; then + echo "--------------------------------------------------------------------------------" + echo "❌ [CONVERSION_FAILURE] Weight conversion logic failed." + echo "Inspect conversion trace in ${log_file}:" + grep -n -C 3 -E "(WeightConverterError|ConversionPlanError|ShapeMismatchError)" "${log_file}" | head -n 30 + echo "--------------------------------------------------------------------------------" + return 1 + elif grep -E "(TPU driver fault|ICI timeout|DMA error|Heartbeat timeout|OutOfMemoryError|SIGKILL|Killed)" "${log_file}"; then + echo "--------------------------------------------------------------------------------" + echo "⚠ī¸ [INFRA_ERROR] Hardware / Driver / Network infrastructure failure detected." + grep -n -C 2 -E "(TPU driver fault|ICI timeout|DMA error|Heartbeat timeout|OutOfMemoryError|SIGKILL|Killed)" "${log_file}" | head -n 20 + echo "--------------------------------------------------------------------------------" + return 101 + elif grep -E "Traceback \(most recent call last\):" "${log_file}"; then + echo "--------------------------------------------------------------------------------" + echo "⚠ī¸ [PYTHON_EXCEPTION] Non-conversion Python exception occurred:" + grep -A 15 "Traceback (most recent call last):" "${log_file}" | tail -n 20 + echo "--------------------------------------------------------------------------------" + return 2 + else + echo "✅ No critical errors detected in recent trainer logs." + echo "Recent output lines:" + tail -n 10 "${log_file}" + return 0 + fi +} + +# ============================================================================== +# Dispatch Command +# ============================================================================== +case "$COMMAND" in + start) + connect_cluster + echo "=================================================================" + echo "Workload ID: ${RUN_ID}" + echo "Model Preset: ${PRESET_NAME} (${MODEL_NAME})" + echo "Image: ${TUNIX_IMAGE}" + echo "Code Sync: ${SYNC_CODE}" + echo "TPU Slices: Trainer=${TRAINER_TPU_SLICE}, Rollout=${ROLLOUT_TPU_SLICE} (x${ROLLOUT_REPLICAS})" + echo "Cluster: ${CLUSTER} (${REGION} / ${PROJECT})" + echo "=================================================================" + + sync_code_to_gcs + stop_workload + start_orchestrator + start_trainer + start_all_rollouts + + if [[ "${DRY_RUN}" != "true" ]]; then + echo "=================================================================" + echo "🎉 Workload launched successfully!" + echo "" + echo "Useful Commands:" + echo " Status: $0 status --run-id=${RUN_ID}" + echo " Logs: $0 logs trainer -f --run-id=${RUN_ID}" + echo " Triage: $0 triage --run-id=${RUN_ID}" + echo " Stop: $0 stop --run-id=${RUN_ID}" + echo "=================================================================" + fi + ;; + + start-trainer) + connect_cluster + echo "Starting trainer only (${TRAINER_ID})..." + start_trainer + ;; + + start-rollout) + connect_cluster + echo "Starting rollouts only (${ROLLOUT_ID})..." + start_all_rollouts + ;; + + render|dry-run) + DRY_RUN="true" + start_orchestrator + start_trainer + start_all_rollouts + ;; + + stop) + connect_cluster + stop_workload + ;; + + restart) + connect_cluster + stop_workload + sleep 3 + "$0" start --run-id="${RUN_ID}" --model="${PRESET_NAME}" --image="${TUNIX_IMAGE}" --namespace="${NAMESPACE:-default}" "${EXTRA_ARGS[@]}" + ;; + + status) + show_status + ;; + + logs) + show_logs "${EXTRA_ARGS[@]}" + ;; + + triage) + triage_workload + ;; + + *) + echo "Error: Unknown command '$COMMAND'." >&2 + print_usage + exit 1 + ;; +esac diff --git a/tunix/experimental/examples/math_gsm8k_dist/launcher.sh b/tunix/experimental/examples/math_gsm8k_dist/launcher.sh index 46138e8f0..17ceacb21 100755 --- a/tunix/experimental/examples/math_gsm8k_dist/launcher.sh +++ b/tunix/experimental/examples/math_gsm8k_dist/launcher.sh @@ -55,6 +55,8 @@ WANDB_RUN_NAME=${WANDB_RUN_NAME:-} WANDB_API_KEY=${WANDB_API_KEY:-} SAMPLER=${SAMPLER:-inprocess_vllm} WEIGHT_SYNC_MODE=${WEIGHT_SYNC_MODE:-none} +MAXTEXT_MODEL_NAME=${MAXTEXT_MODEL_NAME-$(printf '%s' "$MODEL_NAME" | tr '[:upper:]' '[:lower:]')} +MAXTEXT_ATTENTION=${MAXTEXT_ATTENTION:-} PYTHON_BIN=${PYTHON_BIN:-python3} WAIT_TIMEOUT_SECS=${WAIT_TIMEOUT_SECS:-1800} WAIT_POLL_SECS=${WAIT_POLL_SECS:-5} @@ -407,6 +409,14 @@ echo "Launching trainer node on TPU chips $TRAINER_TPU_CHIPS..." --lora_rank="$LORA_RANK" --lora_alpha="$LORA_ALPHA" ) + if [[ -n "$MAXTEXT_CKPT" ]]; then + TRAINER_CMD+=(--maxtext_load_parameters_path="$MAXTEXT_CKPT") + fi + # Same value the rollout gets below, so the two sides cannot disagree about + # which MaxText model they are building. + if [[ -n "$MAXTEXT_MODEL_NAME" ]]; then + TRAINER_CMD+=(--maxtext_model_name="$MAXTEXT_MODEL_NAME") + fi if [[ "$USE_LORA" == "1" || "$USE_LORA" == "true" || "$USE_LORA" == "True" ]]; then TRAINER_CMD+=(--use_lora) fi @@ -456,6 +466,14 @@ echo "Launching rollout node with sampler=$SAMPLER on TPU chips $ROLLOUT_TPU_CHI --lora_alpha="$LORA_ALPHA" --weight_sync_mode="$WEIGHT_SYNC_MODE" ) + # Only the MaxText-native rollout produces tensor names the MaxText trainer + # can match, so weight sync needs this whenever SAMPLER=vllm. + if [[ -n "$MAXTEXT_MODEL_NAME" ]]; then + ROLLOUT_CMD+=( --maxtext_model_name="$MAXTEXT_MODEL_NAME" ) + fi + if [[ -n "$MAXTEXT_ATTENTION" ]]; then + ROLLOUT_CMD+=( --maxtext_attention="$MAXTEXT_ATTENTION" ) + fi if [[ "$USE_LORA" == "1" || "$USE_LORA" == "true" || "$USE_LORA" == "True" ]]; then ROLLOUT_CMD+=(--use_lora) fi diff --git a/tunix/experimental/examples/math_gsm8k_dist/run_gsm8k_dist_grpo.py b/tunix/experimental/examples/math_gsm8k_dist/run_gsm8k_dist_grpo.py index 29f7b3c3d..1a1a0221f 100644 --- a/tunix/experimental/examples/math_gsm8k_dist/run_gsm8k_dist_grpo.py +++ b/tunix/experimental/examples/math_gsm8k_dist/run_gsm8k_dist_grpo.py @@ -150,6 +150,12 @@ def _parse_args(argv: list[str]) -> argparse.Namespace: parser.add_argument("--init_timeout_s", type=float, default=None) parser.add_argument("--inference_addr", type=str, default="") parser.add_argument("--stop_workers_on_exit", action="store_true") + parser.add_argument( + "--rollout_replicas", + type=int, + default=int(os.getenv("ROLLOUT_REPLICAS", "1")), + help="Number of rollout worker replicas to wait for.", + ) parser.add_argument( "--debug", action="store_true", @@ -296,7 +302,7 @@ def main(argv: list[str], context: ProcessContext | None = None) -> None: cluster.wait_for_workers( min_workers={ datatypes.Role.ACTOR: 1, - datatypes.Role.ROLLOUT: 1, + datatypes.Role.ROLLOUT: args.rollout_replicas, datatypes.Role.REFERENCE: 1 if args.beta != 0.0 else 0, }, timeout=args.init_timeout_s, @@ -370,6 +376,9 @@ def main(argv: list[str], context: ProcessContext | None = None) -> None: num_steps=args.max_steps, bring_up=False, ) + except BaseException as exc: + logging.exception("FATAL: StandardRLProgram execution failed: %s", exc) + raise finally: program.close() if args.stop_workers_on_exit: diff --git a/tunix/experimental/examples/math_gsm8k_dist/run_rollout_node.py b/tunix/experimental/examples/math_gsm8k_dist/run_rollout_node.py index d4c051910..1f8617bb4 100644 --- a/tunix/experimental/examples/math_gsm8k_dist/run_rollout_node.py +++ b/tunix/experimental/examples/math_gsm8k_dist/run_rollout_node.py @@ -112,6 +112,12 @@ def _parse_args(argv: list[str]) -> argparse.Namespace: " vllm_batched_rpa)." ), ) + parser.add_argument( + "--maxtext_ckpt_path", + type=str, + default=os.getenv("MAXTEXT_CKPT", ""), + help="Path to MaxText checkpoint to load initial parameters from.", + ) parser.add_argument( "--debug", action="store_true", @@ -127,7 +133,17 @@ def _parse_args(argv: list[str]) -> argparse.Namespace: choices=list(weight_sync_lib.WeightSyncMode), help="Weight sync mode (none, fallback, or raiden).", ) - return parser.parse_args(argv) + parser.add_argument( + "--prefuse_moe_weights", + type=lambda x: str(x).lower() in ("true", "1", "yes"), + default=os.getenv("PREFUSE_MOE_WEIGHTS", "false").lower() in ("true", "1", "yes"), + help="Whether to prefuse MoE weights (gate + up projection).", + ) + parser.add_argument("--tensor_parallel_size", type=int, default=None) + args = parser.parse_args(argv) + if args.tensor_parallel_size is None and (args.sampler_mesh_tp > 1 or args.mesh_tp > 1): + args.tensor_parallel_size = args.sampler_mesh_tp if args.sampler_mesh_tp > 1 else args.mesh_tp + return args def _create_rollout_mesh(args) -> Any: @@ -351,8 +367,8 @@ def _create_vllm_sampler(args): args.maxtext_model_name, ) engine_kwargs["hf_overrides"] = {"architectures": ["MaxTextForCausalLM"]} - # MaxText inference config. prefuse_moe_weights is left False so rollout - # variable names match unfused trainer parameters during weight sync. + # MaxText inference config. When prefuse_moe_weights is True, MoE weights + # are prefused and interleaved per-shard for tensor parallel rollout. maxtext_config_overrides = { "model_name": args.maxtext_model_name, "model_call_mode": "inference", @@ -360,9 +376,14 @@ def _create_vllm_sampler(args): "allow_split_physical_axes": True, "log_config": False, "weight_dtype": "bfloat16", + "prefuse_moe_weights": args.prefuse_moe_weights, } if args.maxtext_attention: maxtext_config_overrides["attention"] = args.maxtext_attention + # Note: load_parameters_path is intentionally NOT set here. + # Checkpoints on disk are scanned (layers.mlp.wi_0...), whereas MaxTextForCausalLM + # is unscanned (layers_0.mlp.wi_0...). Initial weights are pushed via Raiden + # weight sync (_sync_initial_weights) from the trainer after unstacking. engine_kwargs["additional_config"] = { "maxtext_config": maxtext_config_overrides } @@ -389,6 +410,47 @@ def _create_vllm_sampler(args): def main(argv: list[str], context: Any = None) -> None: + try: + import vllm.model_executor.layers.quantization.modelopt as _vllm_modelopt # pylint: disable=g-import-not-at-top + + class _DummyModelOpt: + + def __init__(self, *args, **kwargs): + pass + + def create_weights(self, *args, **kwargs): + pass + + for _name in ( + "ACT", + "WEIGHT", + "CkptCtx", + "KNvfp4Dynamic", + "KNvfp4Static", + "ModelOptNvFp4Config", + "ModelOptNvFp4FusedMoE", + "Shapes", + ): + if not hasattr(_vllm_modelopt, _name): + setattr(_vllm_modelopt, _name, _DummyModelOpt) + except (ImportError, AttributeError) as e: + logging.debug("Optional vllm modelopt patch skipped: %s", e) + + try: + import re # pylint: disable=g-import-not-at-top + import vllm.model_executor.layers.quantization.utils.config_utils as _vllm_config_utils # pylint: disable=g-import-not-at-top + if not hasattr(_vllm_config_utils, "is_equal_or_regex_match"): + def _is_equal_or_regex_match(target_str: str, pattern: str) -> bool: + return target_str == pattern or bool(re.match(pattern, target_str)) + _vllm_config_utils.is_equal_or_regex_match = _is_equal_or_regex_match + except (ImportError, AttributeError) as e: + logging.debug("Optional vllm config_utils patch skipped: %s", e) + + from tunix.experimental.weight_sync.raiden_synchronizer import ( # pylint: disable=g-import-not-at-top + patch_raiden_worker_sync, + ) + patch_raiden_worker_sync() + if context and context.ipc and context.ipc.discovery: pass else: diff --git a/tunix/experimental/examples/math_gsm8k_dist/run_trainer_node.py b/tunix/experimental/examples/math_gsm8k_dist/run_trainer_node.py index bb69e67f5..c5583373b 100644 --- a/tunix/experimental/examples/math_gsm8k_dist/run_trainer_node.py +++ b/tunix/experimental/examples/math_gsm8k_dist/run_trainer_node.py @@ -132,6 +132,18 @@ def _parse_args(argv: list[str]) -> argparse.Namespace: " step 0)." ), ) + parser.add_argument( + "--rollout_mesh_tp", + type=int, + default=0, + help="Rollout tensor parallel mesh dimension for automatic MoE padding calculation.", + ) + parser.add_argument( + "--prefuse_moe_weights", + type=lambda x: str(x).lower() in ("true", "1", "yes"), + default=os.getenv("PREFUSE_MOE_WEIGHTS", "false").lower() in ("true", "1", "yes"), + help="Whether to prefuse MoE weights (gate + up projection).", + ) parser.add_argument( "--debug", action="store_true", @@ -239,6 +251,137 @@ def _load_actor_model(args, mesh: Mesh, *, lora: bool): ) +def _maxtext_modules(): + """Imports MaxText lazily, tolerating both installed layouts. + + Mirrors the guarded import in `tunix/models/automodel.py`: an editable checkout exposes + `maxtext.configs`, while some installs nest it under `maxtext.src.maxtext`. + """ + try: + from maxtext.configs import pyconfig # pylint: disable=g-import-not-at-top + from maxtext.training_engine import maxtext_engine # pylint: disable=g-import-not-at-top + from maxtext.utils import maxtext_utils # pylint: disable=g-import-not-at-top + except ImportError: # pragma: no cover - layout-dependent + from maxtext.src.maxtext.configs import pyconfig # pylint: disable=g-import-not-at-top + from maxtext.src.maxtext.training_engine import maxtext_engine # pylint: disable=g-import-not-at-top + from maxtext.src.maxtext.utils import maxtext_utils # pylint: disable=g-import-not-at-top + return pyconfig, maxtext_engine, maxtext_utils + + +def _tokenizer_pad_id(args) -> int: + """Resolves the pad token id the MaxText adapter masks with. + + Derived exactly as the orchestrator derives it (`run_gsm8k_dist_grpo.py`), including the + eos fallback. The two must agree: the orchestrator pads assembled batches with its id + while the adapter builds `decoder_segment_ids` from this one, so a mismatch silently + corrupts trainer log-probs rather than raising. + """ + from transformers import AutoTokenizer # pylint: disable=g-import-not-at-top + + tokenizer_path = args.tokenizer_path or args.model_dir or args.model_id + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True) + if tokenizer.pad_token_id is None and tokenizer.eos_token is not None: + tokenizer.pad_token = tokenizer.eos_token + return tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0 + + +def _build_maxtext_config(args, num_devices: int) -> Any: + """Builds the MaxText HyperParameters the training engine runs on.""" + pyconfig, _, _ = _maxtext_modules() + + load_parameters_path = args.maxtext_load_parameters_path or args.maxtext_ckpt_path + if not load_parameters_path: + raise ValueError( + "--trainer_backend=maxtext requires --maxtext_load_parameters_path (set " + "MAXTEXT_CKPT). Without it MaxText would random-initialize the model, or try to " + "convert one from HuggingFace." + ) + # MaxText derives micro_batch_size_to_train_on = num_devices * per_device_batch_size, + # and shards the batch dimension of every loss input over the fsdp axis. A microbatch + # that is not a multiple of that axis is rejected at the jit boundary, so reconcile the + # two batch-size sources here rather than letting the failure surface as a sharding + # error deep inside the first step. + if args.train_micro_batch_size % args.mesh_fsdp: + raise ValueError( + f"--train_micro_batch_size={args.train_micro_batch_size} must be a multiple of " + f"--mesh_fsdp={args.mesh_fsdp}; MaxText shards the batch dimension across it." + ) + per_device_batch_size = args.train_micro_batch_size / num_devices + + base_yml = os.path.join( + os.path.dirname(os.path.abspath(pyconfig.__file__)), "base.yml" + ) + if not os.path.exists(base_yml): + raise FileNotFoundError(f"MaxText base.yml not found at {base_yml}") + output_dir = args.maxtext_output_directory or os.path.join(REPO_ROOT, "artifacts", "qwen3_dist_gsm8k", "maxtext") + + argv = [ + "run_trainer_node.py", + base_yml, + f"model_name={args.maxtext_model_name}", + f"run_name={args.worker_id}", + f"base_output_directory={output_dir}", + # load_parameters_path requires enable_checkpointing=True; the config validator + # rejects the combination outright otherwise. + "enable_checkpointing=True", + f"load_parameters_path={load_parameters_path}", + # The trainer stays scanned: the checkpoint is scanned, and the Mode 1 weight-sync + # mapping converts scanned->unscanned for the rollout. + "scan_layers=True", + # Never reach for HuggingFace: with this off and load_parameters_path set, + # from_pretrained cannot take its HF->Orbax conversion path. + "convert_checkpoint_if_possible=False", + # tunix's runtime has already initialized JAX distributed. + "skip_jax_distributed_system=True", + f"per_device_batch_size={per_device_batch_size}", + # tunix owns gradient accumulation: the engine divides by the number of fwd_bwd + # calls it actually saw, so MaxText must not also accumulate. + "gradient_accumulation_steps=1", + f"max_target_length={args.max_prompt_length + args.max_response_length}", + # MaxText's default 'autoselected' attention picks splash attention on TPU, whose + # sa_block_* sizes (512) must divide the sequence length. Here that length is + # max_prompt_length + max_response_length, which the demo sets freely -- 512+128=640 + # already fails, and no fixed block size divides every combination. dot_product has + # no such constraint; at these sequence lengths the difference does not matter. + "attention=dot_product", + f"ici_fsdp_parallelism={args.mesh_fsdp}", + f"ici_tensor_parallelism={args.mesh_tp}", + f"learning_rate={args.learning_rate}", + f"warmup_steps_fraction={args.maxtext_warmup_steps_fraction}", + f"dtype={args.maxtext_dtype}", + f"weight_dtype={args.maxtext_dtype}", + "grad_dtype=float32", + "enable_tensorboard=False", + "record_internal_nn_metrics=False", + "init_weights_seed=42", + f"prefuse_moe_weights={args.prefuse_moe_weights}", + ] + padded_moe_dim = args.maxtext_padded_moe_mlp_dim + if not padded_moe_dim and args.rollout_mesh_tp > 0: + try: + from maxtext.integration.vllm.moe_padding import compute_padded_moe_mlp_dim + tmp_cfg = pyconfig.initialize(argv) + base_dim = getattr(tmp_cfg, "base_moe_mlp_dim", None) or getattr(tmp_cfg, "moe_intermediate_size", None) + if base_dim: + padded_moe_dim = compute_padded_moe_mlp_dim(base_dim, args.rollout_mesh_tp) + logging.info("Auto-computed padded_base_moe_mlp_dim=%d for rollout_mesh_tp=%d", padded_moe_dim, args.rollout_mesh_tp) + except Exception as e: + logging.warning("Could not auto-compute padded_base_moe_mlp_dim: %s", e) + + if padded_moe_dim: + argv.append(f"padded_base_moe_mlp_dim={padded_moe_dim}") + + logging.info("MaxText config argv: %s", argv) + return pyconfig.initialize(argv) + + +def _create_maxtext_mesh(maxtext_config) -> Mesh: + """Builds the mesh MaxText's own sharding annotations are written against.""" + _, _, maxtext_utils = _maxtext_modules() + devices = maxtext_utils.create_device_mesh(maxtext_config) + return Mesh(devices, maxtext_config.mesh_axes) + + class _MeshBoundTrainer: """Binds generic PeftTrainer v2 calls to this worker's JAX mesh.""" @@ -276,6 +419,9 @@ def prepare_weight_sync(self, **kwargs) -> Any: 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"): + logging.info("Checkpoint saving disabled by DISABLE_CHECKPOINTING env var.") + return with self._mesh: self._trainer.save_checkpoint(metadata, **kwargs) @@ -305,6 +451,8 @@ def _create_maxtext_trainer_factory(args) -> Any: load_parameters_path=args.maxtext_ckpt_path, padded_moe_mlp_dim=args.maxtext_padded_moe_mlp_dim, base_output_directory=args.maxtext_output_directory, + rollout_mesh_tp=args.rollout_mesh_tp, + prefuse_moe_weights=args.prefuse_moe_weights, ) logging.info("Creating MaxText device mesh...") mesh = maxtext_utils.create_maxtext_mesh(maxtext_config) diff --git a/tunix/experimental/orchestrator/distributed_rl_engine.py b/tunix/experimental/orchestrator/distributed_rl_engine.py index 49d2529c0..6ba2631cf 100644 --- a/tunix/experimental/orchestrator/distributed_rl_engine.py +++ b/tunix/experimental/orchestrator/distributed_rl_engine.py @@ -22,6 +22,7 @@ import asyncio import collections from collections.abc import Mapping, Sequence +from concurrent import futures import inspect from typing import Any import uuid @@ -113,6 +114,23 @@ def _response_to_trajectory_item(resp: Any) -> datatypes.TrajectoryItem: ) +def _submit_worker( + worker: remote_execution.ActorHandle, + method_name: str, + *args: Any, + **kwargs: Any, +) -> Any: + """Invokes worker.submit safely even when an event loop is active on the current thread.""" + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = None + if loop is not None and loop.is_running(): + with futures.ThreadPoolExecutor(max_workers=1) as pool: + return pool.submit(worker.submit, method_name, *args, **kwargs).result() + return worker.submit(method_name, *args, **kwargs) + + class DistributedRLEngine(rl_engine_interface.AbstractRLEngine): """Worker-backed compute router dispatching RPCs across role pools.""" @@ -325,41 +343,53 @@ async def dispatch_rollouts( return await self.dispatch_rollout_requests(rollout_reqs) async def poll_rollouts( - self, timeout_s: float = remote_execution.LONG_POLL_TIMEOUT_S + self, timeout_s: float = 1.0 ) -> list[datatypes.TrajectoryItem]: """Concurrently long-polls completed rollout responses across all workers.""" if not self._rollout_workers: return [] - async def _poll_worker(worker: remote_execution.ActorHandle) -> Any: + async def _poll_worker(worker: remote_execution.ActorHandle) -> list[Any]: res = worker.poll_responses(timeout_s=timeout_s) if inspect.isawaitable(res): - return await res - return res + res = await res + if res is None or isinstance(res, Exception): + return [] + worker_items = [res] + while True: + nxt = worker.poll_responses(timeout_s=0.0) + if inspect.isawaitable(nxt): + nxt = await nxt + if nxt is None or isinstance(nxt, Exception): + break + worker_items.append(nxt) + return worker_items tasks = [_poll_worker(w) for w in self._rollout_workers] - responses = await asyncio.gather(*tasks, return_exceptions=True) + responses_nested = await asyncio.gather(*tasks, return_exceptions=True) completed: list[datatypes.TrajectoryItem] = [] - for resp in responses: + for resp in responses_nested: if isinstance(resp, Exception) or resp is None: continue - unwrap_fn = getattr(resp, "unwrap", None) - res = ( - unwrap_fn() if callable(unwrap_fn) else getattr(resp, "result", resp) - ) - if res is not None: - items = res if isinstance(res, list) else [res] - for it in items: - if isinstance(it, dict): - it = datatypes.RolloutResponse(**it) - traj_item = _response_to_trajectory_item(it) - logging.debug( - "Received rollout response (prompt_id=%s, group_index=%d).", - traj_item.prompt_id, - traj_item.group_index, - ) - completed.append(traj_item) + items_list = resp if isinstance(resp, list) else [resp] + for r in items_list: + unwrap_fn = getattr(r, "unwrap", None) + res = ( + unwrap_fn() if callable(unwrap_fn) else getattr(r, "result", r) + ) + if res is not None: + items = res if isinstance(res, list) else [res] + for it in items: + if isinstance(it, dict): + it = datatypes.RolloutResponse(**it) + traj_item = _response_to_trajectory_item(it) + logging.debug( + "Received rollout response (prompt_id=%s, group_index=%d).", + traj_item.prompt_id, + traj_item.group_index, + ) + completed.append(traj_item) return completed async def generate( @@ -576,14 +606,14 @@ def configure_worker( "Auto-configuring trainer loss and model input fn on %s worker...", role_name, ) - worker.submit("with_loss_fn", algo.loss_fn(), has_aux=True) + _submit_worker(worker, "with_loss_fn", algo.loss_fn(), has_aux=True) pad_id = getattr(assembler, "pad_id", kwargs.get("pad_id", 0)) eos_id = getattr(assembler, "eos_id", kwargs.get("eos_id", pad_id)) gen_fn = algo.build_gen_model_input_fn( pad_id=pad_id, # pyrefly: ignore[bad-argument-type] eos_id=eos_id, # pyrefly: ignore[bad-argument-type] ) - worker.submit("with_gen_model_input_fn", gen_fn) + _submit_worker(worker, "with_gen_model_input_fn", gen_fn) case datatypes.Role.ROLLOUT: if not self._rollout_workers: diff --git a/tunix/experimental/orchestrator/rl_engine_interface.py b/tunix/experimental/orchestrator/rl_engine_interface.py index 3dea4a4c4..5332dc193 100644 --- a/tunix/experimental/orchestrator/rl_engine_interface.py +++ b/tunix/experimental/orchestrator/rl_engine_interface.py @@ -83,7 +83,7 @@ async def dispatch_rollouts( ... async def poll_rollouts( - self, timeout_s: float = remote_execution.LONG_POLL_TIMEOUT_S + self, timeout_s: float = 1.0 ) -> list[datatypes.TrajectoryItem]: """Retrieves completed rollout responses from workers via long-polling.""" ... diff --git a/tunix/experimental/orchestrator/rl_program.py b/tunix/experimental/orchestrator/rl_program.py index 6e3717889..a0206c50a 100644 --- a/tunix/experimental/orchestrator/rl_program.py +++ b/tunix/experimental/orchestrator/rl_program.py @@ -22,6 +22,7 @@ import asyncio from collections.abc import Callable, Iterable, Sequence import dataclasses +import os import time from typing import Any @@ -35,6 +36,13 @@ from tunix.experimental.queue_manager import trajectory_queue_manager from tunix.sft import metrics_logger as metrics_logger_lib +# Bounds on the pre-dispatch weight sync. The wait is dominated by vLLM +# EngineCore startup (~20s locally, longer once weights come off GCS), so the +# timeout is generous; failing here is better than silently rolling out from +# uninitialized weights. +_INITIAL_SYNC_TIMEOUT_S = 600.0 +_INITIAL_SYNC_RETRY_S = 5.0 + MetricsLogger = metrics_logger_lib.MetricsLogger MetricsLoggerOptions = metrics_logger_lib.MetricsLoggerOptions Mode = metrics_logger_lib.Mode @@ -714,15 +722,18 @@ async def train_stage(self) -> None: # TODO(tunix-dev): For now any failures in save_checkpoint will # abort the entire program. Make it configurable on whether to fail # or continue. - await self.engine.save_checkpoint( - role=datatypes.Role.ACTOR, - metadata={ - "step": self.step + 1, - "policy_version": self.policy_version, - "num_rollouts": num_rollouts, - "num_microbatches": num_microbatches, - }, - ) + if os.environ.get("DISABLE_CHECKPOINTING", "false").lower() in ("true", "1"): + logging.info("Skipping save_checkpoint as DISABLE_CHECKPOINTING is set.") + else: + await self.engine.save_checkpoint( + role=datatypes.Role.ACTOR, + metadata={ + "step": self.step + 1, + "policy_version": self.policy_version, + "num_rollouts": num_rollouts, + "num_microbatches": num_microbatches, + }, + ) if not scored_items: break @@ -792,6 +803,56 @@ async def train_stage(self) -> None: self.on_step_end(current_step, step_result) self._step += 1 + 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 + built for itself is what generates the step-0 trajectories. That is fine + when the sampler loads the same checkpoint as the trainer, but the MaxText + -in-vLLM path does not: `MaxTextForCausalLM` is constructed with an empty + `load_parameters_path`, so `from_pretrained` skips the Orbax restore and + leaves the model randomly initialized. The resulting rollouts are + degenerate (one token repeated to the length cap) yet still get scored and + trained on, so the run looks healthy while learning from noise. + + Syncing once up front makes policy version 0 mean "the trainer's starting + weights" on both sides. The retry loop exists because bringing up workers + does not wait for readiness -- the vLLM EngineCore spawns as a separate + process and takes ~20s to come up, well after this coroutine starts. + """ + if not self.sync_weights: + return + assert self.engine is not None + + deadline = time.monotonic() + _INITIAL_SYNC_TIMEOUT_S + attempt = 0 + while True: + attempt += 1 + try: + version = await self.engine.sync_weights(role=datatypes.Role.ACTOR) + except Exception as exc: # pylint: disable=broad-except + if time.monotonic() >= deadline: + raise RuntimeError( + "Initial weight sync never succeeded; rollout workers would" + " generate from uninitialized weights." + ) from exc + logging.info( + "Initial weight sync attempt %d not ready yet (%s); retrying.", + attempt, + exc, + ) + await asyncio.sleep(_INITIAL_SYNC_RETRY_S) + continue + # Track the engine's counter rather than staying at 0, so the version + # the first rollouts are tagged with is the one they actually ran. + if version is not None: + self.policy_version = version + logging.info( + "Initial weight sync complete; rollout policy_version=%d.", + self.policy_version, + ) + return + async def run_async( self, engine: rl_engine_interface.AbstractRLEngine, @@ -811,13 +872,15 @@ async def run_async( if self.sync_weights: await engine.prepare_rollout_policy( role=datatypes.Role.ACTOR, - sync_weights=True, + sync_weights=False, policy_version=self.policy_version, ) max_groups_ahead = self.mini_batch_size * (self.max_staleness + 1) self._dispatch_capacity = asyncio.Semaphore(max_groups_ahead) + await self._sync_initial_weights() + train_task = asyncio.create_task(self.train_stage()) tasks = [ asyncio.create_task(self.rollout_dispatch_stage()), diff --git a/tunix/experimental/rollout/collector.py b/tunix/experimental/rollout/collector.py index 115657bb1..3d7f276f7 100644 --- a/tunix/experimental/rollout/collector.py +++ b/tunix/experimental/rollout/collector.py @@ -14,6 +14,7 @@ """Trajectory Collector Engine wrapping TrajectoryCollectEngine with pause/resume/cancel control.""" +import logging from typing import Any, List import numpy as np from tunix.experimental.common import datatypes @@ -67,6 +68,7 @@ def __init__( self.max_response_length = request.generation_kwargs.get( "max_response_length" ) + self._accumulated_token_ids: List[int] = [] async def run_episode(self) -> trajectory_lib.Trajectory: """Executes multi-turn agentic rollout episode and returns standardized Trajectory.""" @@ -107,6 +109,25 @@ async def model_call( else: prompt_tokens = np.array([[0]], dtype=np.int32) + # TEMPORARY instrumentation: nothing else on this path records what the + # sampler actually produced, so a run of empty completions reports the + # same "rollouts=N" summary as a real one. Remove once the vllm sampler + # path is trusted. + logging.info( + "[collector] traj=%s completion_tokens=%d prompt_tokens=%d" + " logprobs=%s text=%r", + self.traj_id, + np.asarray(tokens).size, + prompt_tokens.size, + "none" if logprobs is None else np.asarray(logprobs).size, + text[:160], + ) + + if not self._accumulated_token_ids and prompt_tokens.size: + self._accumulated_token_ids.extend([int(x) for x in prompt_tokens.reshape(-1)]) + if np.asarray(tokens).size: + self._accumulated_token_ids.extend([int(x) for x in np.asarray(tokens).reshape(-1)]) + return base_rollout.RolloutOutput( text=[text], logits=None, @@ -200,4 +221,8 @@ def cancel(self) -> None: def get_accumulated_token_ids(self) -> List[int]: """Returns token IDs of historical turns for Raiden KV-cache transfer.""" + if self._accumulated_token_ids: + return list(self._accumulated_token_ids) + if hasattr(self.request, "prompt_tokens") and self.request.prompt_tokens: + return list(self.request.prompt_tokens) return [] diff --git a/tunix/experimental/rollout/inprocess_vllm_sampler_adapter.py b/tunix/experimental/rollout/inprocess_vllm_sampler_adapter.py index 3a0260c17..f99d3434e 100644 --- a/tunix/experimental/rollout/inprocess_vllm_sampler_adapter.py +++ b/tunix/experimental/rollout/inprocess_vllm_sampler_adapter.py @@ -59,7 +59,7 @@ def __init__( elif isinstance(weight_sync_mode, str): self.weight_sync_mode = weight_sync.WeightSyncMode(weight_sync_mode) else: - self.weight_sync_mode = weight_sync.WeightSyncMode.FALLBACK + self.weight_sync_mode = weight_sync.DEFAULT_WEIGHT_SYNC_MODE self.enable_raiden = ( self.weight_sync_mode == weight_sync.WeightSyncMode.RAIDEN ) @@ -75,7 +75,9 @@ def __init__( from tunix.experimental.weight_sync import raiden_weight_sync_delegate # pylint: disable=g-import-not-at-top self.raiden_sync_delegate = ( - raiden_weight_sync_delegate.RaidenWeightSyncDelegate() + raiden_weight_sync_delegate.RaidenWeightSyncDelegate( + server_id=self.server_id + ) ) if not self.enable_raiden and self.raiden_sync_delegate: @@ -454,6 +456,19 @@ async def post_weight_sync( ) return True + async def abort_weight_sync( + self, + sync_request: base_sampler_lib.WeightSyncRequest | Any = None, + **kwargs, + ) -> str | None | Any: + """Safely aborts weight sync round.""" + if self.enable_raiden and self.raiden_sync_delegate: + if hasattr(self.raiden_sync_delegate, "abort_weight_sync"): + return await self.raiden_sync_delegate.abort_weight_sync( + sync_request=sync_request, **kwargs + ) + return True + async def get_transfer_status(self, req_id: str | Any, **kwargs) -> str | Any: """Queries status of an ongoing weight transfer or KV-cache migration.""" del req_id, kwargs diff --git a/tunix/experimental/rollout/manager.py b/tunix/experimental/rollout/manager.py index ed7ce7726..36b9c306e 100644 --- a/tunix/experimental/rollout/manager.py +++ b/tunix/experimental/rollout/manager.py @@ -66,7 +66,7 @@ def __init__( if sampler is None: sampler_type = getattr(config, "sampler_type", "vanilla") weight_sync_mode = getattr( - config, "weight_sync_mode", weight_sync.WeightSyncMode.FALLBACK + config, "weight_sync_mode", weight_sync.DEFAULT_WEIGHT_SYNC_MODE ) if sampler_type == "vllm": @@ -85,7 +85,9 @@ def __init__( from tunix.experimental.weight_sync import raiden_weight_sync_delegate # pylint: disable=g-import-not-at-top raiden_delegate = ( - raiden_weight_sync_delegate.RaidenWeightSyncDelegate() + raiden_weight_sync_delegate.RaidenWeightSyncDelegate( + server_id="inprocess_vllm_sampler" + ) ) sampler = inprocess_vllm_sampler_adapter.InprocessVllmSamplerAdapter( # pyrefly: ignore[bad-instantiation] @@ -101,7 +103,9 @@ def __init__( from tunix.experimental.weight_sync import raiden_weight_sync_delegate # pylint: disable=g-import-not-at-top raiden_delegate = ( - raiden_weight_sync_delegate.RaidenWeightSyncDelegate() + raiden_weight_sync_delegate.RaidenWeightSyncDelegate( + server_id="vanilla_sampler" + ) ) sampler = vanilla_sampler_adapter.VanillaSamplerAdapter( @@ -337,6 +341,17 @@ async def post_weight_sync( self._traffic.reopen() return res + async def abort_weight_sync( + self, sync_request: sampler_lib.WeightSyncRequest | Any = None, **kwargs + ) -> Any: + """Discards the round, delegates to sampler if available, and resumes serving.""" + res = None + if self.sampler and hasattr(self.sampler, "abort_weight_sync"): + res = await self.sampler.abort_weight_sync(sync_request, **kwargs) + self.resume_all() + self.reopen_admission() + return res + def reopen_admission(self) -> bool: """Reopens rollout admission after an aborted round.""" return self._traffic.reopen() diff --git a/tunix/experimental/rollout/sampler.py b/tunix/experimental/rollout/sampler.py index b8927ab43..09c799e89 100644 --- a/tunix/experimental/rollout/sampler.py +++ b/tunix/experimental/rollout/sampler.py @@ -113,25 +113,8 @@ def __post_init__(self): ) -@dataclasses.dataclass(kw_only=True) -class WeightSyncRequest(datatypes.Request): - """Configuration and routing metadata for synchronizing policy model weights. - - Attributes: - controller_id: Optional identifier for transport controllers (e.g., TPU - Raiden). - policy_version: Target policy version identifier of the weights to sync. - weights: Optional source weights payload for non-Raiden / fallback sync. - source_metadata: Optional transport/layout metadata describing source - weights. - extra_config: Optional backend-specific configuration parameters. - """ - - controller_id: str = "" - policy_version: int = 0 - weights: Any = None - source_metadata: Any = None - extra_config: dict[str, Any] = dataclasses.field(default_factory=dict) +# Alias canonical DTO from datatypes module for backwards compatibility. +WeightSyncRequest = datatypes.WeightSyncRequest @dataclasses.dataclass(kw_only=True) @@ -220,6 +203,12 @@ async def post_weight_sync( """Finalizes and switches active policy weights after transfer completion.""" ... + async def abort_weight_sync( + self, sync_request: WeightSyncRequest | Any = None, **kwargs + ) -> str | None | Any: + """Discards staging and restores serving the previous policy weights.""" + ... + async def get_transfer_status( self, req_id: str | Any, **kwargs ) -> str | Any: diff --git a/tunix/experimental/rollout/vanilla_sampler_adapter.py b/tunix/experimental/rollout/vanilla_sampler_adapter.py index ec8c9a948..ffa9662a1 100644 --- a/tunix/experimental/rollout/vanilla_sampler_adapter.py +++ b/tunix/experimental/rollout/vanilla_sampler_adapter.py @@ -65,7 +65,7 @@ def __init__( self.config = config self.raiden_sync_delegate = raiden_sync_delegate self.weight_sync_mode = getattr( - config, "weight_sync_mode", weight_sync.WeightSyncMode.FALLBACK + config, "weight_sync_mode", weight_sync.DEFAULT_WEIGHT_SYNC_MODE ) self.enable_raiden = ( self.weight_sync_mode == weight_sync.WeightSyncMode.RAIDEN @@ -75,7 +75,9 @@ def __init__( from tunix.experimental.weight_sync import raiden_weight_sync_delegate # pylint: disable=g-import-not-at-top self.raiden_sync_delegate = ( - raiden_weight_sync_delegate.RaidenWeightSyncDelegate() + raiden_weight_sync_delegate.RaidenWeightSyncDelegate( + server_id=self.server_id + ) ) if not self.enable_raiden and self.raiden_sync_delegate: @@ -233,7 +235,8 @@ async def sample( if hasattr(req, "sampling_params") and req.sampling_params is not None else base_sampler_lib.SamplingParams() ) - assert sp is not None + if sp is None: + raise ValueError("SamplingParams cannot be None") max_gen_steps_list.append(sp.max_tokens) temps.append(sp.temperature) @@ -471,6 +474,19 @@ async def post_weight_sync( # Fallback mode: acts as a no-op returning True. return True + async def abort_weight_sync( + self, + sync_request: base_sampler_lib.WeightSyncRequest | Any = None, + **kwargs, + ) -> str | None | Any: + """Safely aborts weight sync round.""" + if self.enable_raiden and self.raiden_sync_delegate: + if hasattr(self.raiden_sync_delegate, "abort_weight_sync"): + return await self.raiden_sync_delegate.abort_weight_sync( + sync_request=sync_request, **kwargs + ) + return True + # --- KV-cache Migration --- async def migrate_kv_cache( self, diff --git a/tunix/experimental/rollout/vllm_sampler_adapter.py b/tunix/experimental/rollout/vllm_sampler_adapter.py index a6e51f8d5..a8643dc54 100644 --- a/tunix/experimental/rollout/vllm_sampler_adapter.py +++ b/tunix/experimental/rollout/vllm_sampler_adapter.py @@ -94,6 +94,7 @@ def __init__( model_name: str = "", sampler_instance: Any = None, worker_index: int = 0, + raiden_job_name: str = "", parallelism: int = 4, weight_sync_mode: weight_sync.WeightSyncMode | str | None = None, **kwargs, @@ -103,6 +104,12 @@ def __init__( self.model_name = model_name or (engine_args.model if engine_args else "") self.sampler = sampler_instance self.worker_index = worker_index + # Raiden treats units sharing a job_name as hosts of ONE job and splits the + # weights across them (`num_dst_physical_hosts` in raiden_controller), so + # independent rollout replicas each need their own job_name to be sent a + # full copy. server_id already has that granularity; worker_index stays the + # host index *within* one replica. + self.raiden_job_name = raiden_job_name or self.server_id self._parallelism = parallelism # Defaults to RAIDEN when unspecified: RLVllmSampler drives weight sync @@ -184,6 +191,34 @@ async def start(self, **kwargs) -> Any: """Starts the underlying sampler engine.""" return await self._require_sampler().start(**kwargs) + async def _ensure_started(self) -> None: + """Brings the engine up if nothing has needed it yet. + + RLVllmSampler builds its AsyncLLM lazily and `sample()` is the only caller + of `start()`. Weight sync needs the engine too -- it owns the TPU worker, + and therefore the Raiden binding -- and the first sync lands before the + first sample, because `prepare_rollout_policy` syncs ahead of dispatch. + Without this the round finds no worker, reports an empty destination + manifest, and deadlocks: the engine waits for a sample that dispatch is + waiting on the sync to allow. Guarded on `_is_running` rather than calling + `start()` unconditionally, because `start()` warns when the engine is + already up and this runs on every sync round. + + TODO(tunix-dev): drop this once the orchestrator owns rollout-worker + lifecycle and can guarantee the engine is up before it issues any phase + call; the ordering belongs there, not in a guard on each entry point. + """ + if self.sampler is None: + self.initialize() + sampler = self._require_sampler() + if not getattr(sampler, "_is_running", False): + logger.info( + "VllmSamplerAdapter [%s] starting engine for weight sync (no" + " sample has forced it up yet).", + self.server_id, + ) + await sampler.start() + async def stop(self, **kwargs) -> Any: """Stops the underlying sampler engine.""" return await self._require_sampler().stop(**kwargs) @@ -243,8 +278,11 @@ async def bind_weight_sync( del sync_request, kwargs if not self.enable_raiden: return None + await self._ensure_started() return await self._require_sampler().bind_raiden_sync( - worker_index=self.worker_index, parallelism=self._parallelism + worker_index=self.worker_index, + parallelism=self._parallelism, + job_name=self.raiden_job_name, ) async def get_weight_sync_metadata( @@ -259,6 +297,7 @@ async def get_weight_sync_metadata( " get_weight_sync_metadata when Raiden is disabled" f" (weight_sync_mode={self.weight_sync_mode.value})." ) + await self._ensure_started() meta = await self._require_sampler().get_raiden_metadata() return [weight_sync.WorkUnitMetadata.from_dict(m) for m in meta or []] @@ -306,6 +345,16 @@ async def weight_sync(self, sync_request: Any = None, **kwargs: Any) -> Any: result = sampler.refresh_model_state_leaves() if asyncio.iscoroutine(result): await result + else: + # Never silently skip this: the sampler's runner dispatches through a + # `state_leaves` view derived from `state` at load time, so without the + # refresh it keeps serving pre-sync weights and the only symptom is + # garbage completions. + logger.warning( + "sampler %s has no refresh_model_state_leaves(); the rollout's" + " state_leaves are not re-pointed after h2d.", + type(sampler).__name__, + ) self._tracker.complete(sync_request, "h2d_done") return True @@ -382,6 +431,13 @@ async def get_transfer_status(self, req_id: Any, **kwargs) -> Any: return await sampler.get_transfer_status(req_id, **kwargs) return "UNKNOWN" + def get_target_state(self) -> Any: + """Returns target state shape/dtype pytree for weight conversion.""" + sampler = self._require_sampler() + if hasattr(sampler, "get_target_state"): + return sampler.get_target_state() + return None + async def get_load_info(self, **kwargs) -> base_sampler_lib.LoadInfo: """Returns load information from the underlying engine.""" info = await self._require_sampler().get_load_info(**kwargs) diff --git a/tunix/experimental/train/peft_trainer_v2.py b/tunix/experimental/train/peft_trainer_v2.py index 38e082e0c..b72e026a3 100644 --- a/tunix/experimental/train/peft_trainer_v2.py +++ b/tunix/experimental/train/peft_trainer_v2.py @@ -1194,7 +1194,8 @@ def prepare_weight_sync(self, sync_request: Any = None, **kwargs) -> Any: mapping_config = mappings_lib.MappingConfig.build( model=self.model, backend=backend ) - except Exception: # pylint: disable=broad-exception-caught + except (ImportError, AttributeError, ValueError) as e: + logging.warning("Failed to build mapping_config for backend %s: %s", backend, e) mapping_config = None if ( diff --git a/tunix/experimental/weight_sync/raiden_handler.py b/tunix/experimental/weight_sync/raiden_handler.py index 6e6706e18..a1f4d0ba3 100644 --- a/tunix/experimental/weight_sync/raiden_handler.py +++ b/tunix/experimental/weight_sync/raiden_handler.py @@ -35,6 +35,7 @@ import asyncio import dataclasses import logging +import math import threading from typing import Any, Mapping, Optional, Sequence @@ -273,16 +274,23 @@ def _validate_metadata(metadata: weight_sync.WorkUnitMetadata) -> None: f" of {tensor.name!r} must have logical mesh size 1, got" f" {logical_size}" ) - elif axis not in physical_axes: + continue + # A dimension may be sharded over the product of several axes, which + # the wire form spells comma-joined; the controller splits it the same + # way to fold the sub-axis coordinates into one tensor coordinate. + sub_axes = axis.split(",") + unknown = [a for a in sub_axes if a not in physical_axes] + if unknown: raise ValueError( f"work unit {metadata.unit}: variable {tensor.name!r} names" f" unknown mesh axis {axis!r}" ) - elif logical_size != physical_axes[axis]: + physical_size = math.prod(physical_axes[a] for a in sub_axes) + if logical_size != physical_size: raise ValueError( f"work unit {metadata.unit}: variable {tensor.name!r} maps axis" f" {axis!r} to logical size {logical_size}, but the physical" - f" mesh has size {physical_axes[axis]}" + f" mesh has size {physical_size}" ) def transfer( diff --git a/tunix/experimental/weight_sync/raiden_synchronizer.py b/tunix/experimental/weight_sync/raiden_synchronizer.py index 955250a67..56602cb3d 100644 --- a/tunix/experimental/weight_sync/raiden_synchronizer.py +++ b/tunix/experimental/weight_sync/raiden_synchronizer.py @@ -17,6 +17,8 @@ from __future__ import annotations import collections +import dataclasses +import gc import inspect import ipaddress import os @@ -25,6 +27,7 @@ from absl import logging import jax +from jax.experimental import compute_on import jax.numpy as jnp from tunix.experimental.weight_sync import weight_sync @@ -47,17 +50,40 @@ def _log_rss(tag: str) -> None: ) -_ws_lib: Any = None -try: - from tpu_sync.api.jax import weight_synchronizer as _ws_lib # pytype: disable=import-error pylint: disable=g-import-not-at-top -except ImportError: - _ws_lib = None +_lazy_modules: dict[str, Any] = {} -_raiden_ffi: Any = None -try: - from tpu_sync.frameworks.jax import weight_synchronizer_ffi as _raiden_ffi # pytype: disable=import-error pylint: disable=g-import-not-at-top -except ImportError: - _raiden_ffi = None + +def _lazy_import_module(module_path: str) -> Any: + """Imports a module lazily by path, caching the result.""" + if module_path not in _lazy_modules: + try: + import importlib # pylint: disable=g-import-not-at-top + _lazy_modules[module_path] = importlib.import_module(module_path) + except ImportError: + _lazy_modules[module_path] = None + return _lazy_modules[module_path] + + +def _get_ws_lib() -> Any: + """Imports tpu_sync weight_synchronizer lazily to prevent early C++ library symbol collisions.""" + if "_ws_lib" in globals(): + return globals()["_ws_lib"] + return _lazy_import_module("tpu_sync.api.jax.weight_synchronizer") + + +def _get_raiden_ffi() -> Any: + """Imports tpu_sync weight_synchronizer_ffi lazily to avoid loading XLA runtime early.""" + if "_raiden_ffi" in globals(): + return globals()["_raiden_ffi"] + return _lazy_import_module("tpu_sync.frameworks.jax.weight_synchronizer_ffi") + + +def __getattr__(name: str) -> Any: + if name == "_raiden_ffi": + return _get_raiden_ffi() + if name == "_ws_lib": + return _get_ws_lib() + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") def _ensure_ffi_compute_on_compat() -> None: @@ -88,9 +114,20 @@ def _ensure_ffi_compute_on_compat() -> None: ) compute_on_mod.compute_on = compute_on2 + # The wheel does `from jax.experimental import compute_on` and then calls + # `compute_on.compute_on(...)`. jax.experimental.compute_on binds the name at + # import time, so patching jax._src alone leaves the caller on the old + # two-arg version and the decorator dies with + # TypeError: compute_on() got an unexpected keyword argument 'out_memory_spaces' + try: + from jax.experimental import compute_on as _public_compute_on # pytype: disable=import-error pylint: disable=g-import-not-at-top + + _public_compute_on.compute_on = compute_on2 + except ImportError: + pass logging.warning( - "Patched jax._src.compute_on.compute_on to compute_on2 for TPU-sync FFI" - " compatibility." + "Patched jax._src.compute_on.compute_on (and jax.experimental." + "compute_on) to compute_on2 for TPU-sync FFI compatibility." ) @@ -123,6 +160,32 @@ def unpack_ip(row: Any) -> str: return f"[{addr_str}]" if ":" in addr_str else addr_str +def to_host_cpu_state(state: Any) -> Any: + """Copies arrays to client host memory; proxy arrays cannot bind directly. + + The legacy transport needs real host buffers, which is what lets a Pathways + trainer feed a destination that has no FFI handlers (an mcjax sampler). + """ + cpu = jax.local_devices(backend="cpu")[0] + leaves, treedef = jax.tree_util.tree_flatten(state) + new_leaves = [] + for i, leaf in enumerate(leaves): + leaves[i] = None # drop our ref as we go, so the copy peaks at one model + arr = getattr(leaf, "value", leaf) + # Single hop: device_get + device_put would make two host copies. + new_leaves.append( + jax.device_put(arr, cpu) + if hasattr(arr, "shape") and hasattr(arr, "dtype") + else leaf + ) + # Proxy transit buffers land in reference cycles, so refcounting alone + # does not reclaim them between leaves. + if i % 4 == 3: + gc.collect() + gc.collect() + return jax.tree_util.tree_unflatten(treedef, new_leaves) + + def flatten_weights(state: Any) -> Tuple[List[str], List[Any]]: """Returns (names, arrays) for every array leaf, in stable tree order.""" names, arrays = [], [] @@ -134,6 +197,27 @@ def flatten_weights(state: Any) -> Tuple[List[str], List[Any]]: return names, arrays +def _normalize_param_name(name: str) -> str: + for prefix in ("['base']", "['model']", "base.", "model."): + if name.startswith(prefix): + name = name[len(prefix):] + if name.endswith(".value"): + name = name[:-len(".value")] + return name + + +def _canonicalize_param_name(name: str) -> str: + norm = _normalize_param_name(name) + import re + m = re.match(r"^\['layers'\]\['(\d+)'\](.*)$", norm) + if m: + return f"['layers_{m.group(1)}']{m.group(2)}" + m = re.match(r"^layers\.(\d+)\.(.*)$", norm) + if m: + return f"layers_{m.group(1)}.{m.group(2)}" + return norm + + def _bindable(arr: Any, *, allow_proxy: bool = False) -> bool: """True if the native layer can bind this leaf. @@ -196,14 +280,74 @@ def _axis_name(axis: Any) -> str: return ",".join(axis) +def _devices_per_host(devices: List[Any]) -> int: + """Devices sharing one physical host, i.e. Raiden's `num_shards`. + + The native layer derives `submanager_idx = shard_idx / num_shards` and + `slot = shard_idx % num_shards`, so this must be the real per-host device + count. Overstate it and every host allocates staging for the whole slice but + fills only its own share, leaving the rest of its SetGlobalShardIndices at + -1 -- the transfer then completes green while delivering only the shards one + host happened to own. + + `process_index` alone is wrong under Pathways: the client is a single process + driving every worker, so all proxy devices report 0 and this collapses to + len(devices). Prefer whichever attribute actually distinguishes the workers, + and fall back to the per-host hardware ordinal. + """ + env = os.environ.get("RAIDEN_DEVICES_PER_HOST") + if env: + n = int(env) + if n > 0 and len(devices) % n == 0: + return n + logging.warning( + "ignoring RAIDEN_DEVICES_PER_HOST=%s: not a divisor of %d devices", + env, len(devices)) + for attr in ("task_id", "process_index"): + groups = {getattr(d, attr, None) for d in devices} + groups.discard(None) + if len(groups) > 1 and len(devices) % len(groups) == 0: + return len(devices) // len(groups) + local_ids = {getattr(d, "local_hardware_id", None) for d in devices} + local_ids.discard(None) + if 1 < len(local_ids) < len(devices) and len(devices) % len(local_ids) == 0: + return len(local_ids) + return len(devices) + + +def _reduce_mesh(mesh: Any) -> Any: + """Drops size-1 axes from a mesh for Raiden's FFI shard_map. + + `init_weight_synchronizer` specs its inputs as `P(*mesh.axis_names)`, so a + MaxText mesh -- twelve axes, most of them singletons -- yields a spec longer + than any operand's rank and shard_map rejects it. Only the trivial axes go; + the real sharding is preserved. Mirrors tpu-inference's `_reduce_mesh`. + """ + keep = [a for a in mesh.axis_names if int(mesh.shape[a]) > 1] + if not keep or len(keep) == len(mesh.axis_names): + return mesh + return jax.sharding.Mesh( + mesh.devices.reshape(tuple(int(mesh.shape[a]) for a in keep)), + axis_names=tuple(keep), + ) + + def _tensor_metadata(name: str, arr: Any, layer_idx: int): sharding: Any = getattr(arr, "sharding", None) spec = tuple(getattr(sharding, "spec", ()) or ()) spec = (spec + (None,) * arr.ndim)[: arr.ndim] - try: - local = sharding.shard_shape(tuple(arr.shape)) - mesh_shape = tuple(g // l for g, l in zip(arr.shape, local)) - except Exception: # pylint: disable=broad-exception-caught + if sharding is not None and hasattr(sharding, "shard_shape"): + try: + local = sharding.shard_shape(tuple(arr.shape)) + mesh_shape = tuple(g // l for g, l in zip(arr.shape, local)) + except Exception as e: # pylint: disable=broad-exception-caught + logging.warning( + "Could not compute mesh_shape for %s from sharding: %s, falling back to 1D", + name, + e, + ) + mesh_shape = (1,) * arr.ndim + else: mesh_shape = (1,) * arr.ndim return weight_sync.TensorMetadata( name=name, @@ -221,9 +365,8 @@ class RaidenSynchronizer: Used by both the trainer and the sampler. Construct with a state to bind right away, or leave it out and call `bind` when the weights exist; every - later `bind` rebinds the same transport. Known limits: one mesh axis per - tensor dim, and without the tpu_sync wheel the metadata carries no shard - addresses, so the handler refuses registration. + later `bind` rebinds the same transport. Without the tpu_sync wheel the + metadata carries no shard addresses, so the handler refuses registration. """ def __init__( @@ -233,10 +376,32 @@ def __init__( *, worker_index: int = 0, auto_h2d: bool = False, + host_stage: Optional[bool] = None, + use_ffi: Optional[bool] = None, parallelism: int = 4, bind_ip: Optional[str] = None, ): is_proxy = "proxy" in os.environ.get("JAX_PLATFORMS", "") + # Pathways says where the arrays live; FFI says which transport moves them. + # An FFI source only feeds an FFI destination, so a Pathways trainer + # serving an mcjax sampler runs RAIDEN_USE_FFI=0 on both sides and host + # stages instead. Same rule as tpu-inference's use_ffi(): proxy picks the + # default, an explicit env var wins, and no wheel means no FFI. + default_ffi = "1" if is_proxy else "0" + if use_ffi is None: + use_ffi = (os.environ.get("RAIDEN_USE_FFI", default_ffi) == "1" + and _get_raiden_ffi() is not None) + # host_stage and use_ffi are mutually exclusive -- FFI binds the proxy + # arrays in place, and staged CPU arrays carry no mesh for + # _init_ffi_transport -- but they DO arrive in contradiction: the pinned + # MaxText in our image passes host_stage=is_pathways unconditionally + # (site-packages/maxtext/training_engine/maxtext_engine.py:1086). Normalise + # here rather than guarding at each use. Grep the *pinned* MaxText, not the + # maxtext checkout, before concluding these parameters are unused. + if use_ffi: + host_stage = False + elif host_stage is None: + host_stage = is_proxy self.job_name = job_name self.worker_index = worker_index self.names: List[str] = [] @@ -244,10 +409,13 @@ def __init__( self.ip = bind_ip or local_ip() self._auto_h2d = auto_h2d self._is_proxy = is_proxy + self._use_ffi = use_ffi + self._host_stage = host_stage self._parallelism = parallelism self._sync: Any = None self._ips: List[str] = [] self._unique_listeners: List[str] = [] + self._listeners: List[str] = [] self._ffi_mesh: Any = None self._ffi_shard_idx: Any = None if state is not None: @@ -259,6 +427,10 @@ def bound(self) -> bool: @property def active(self) -> bool: + # Under FFI there is no native `_sync` and `_ips` fill only in d2h(), so + # bound arrays are the only pre-d2h signal. MaxText gates d2h() on this. + if self._use_ffi: + return self.bound return self._sync is not None or bool(self._ips) def _init_ffi_transport(self, *, is_d2h: bool) -> None: @@ -266,7 +438,8 @@ def _init_ffi_transport(self, *, is_d2h: bool) -> None: raise RuntimeError( f"{self.job_name}: bind() must stage arrays before FFI init" ) - if _raiden_ffi is None: + raiden_ffi = _get_raiden_ffi() + if raiden_ffi is None: raise RuntimeError( "weight_synchronizer_ffi is not available for FFI weight sync." ) @@ -278,6 +451,7 @@ def _init_ffi_transport(self, *, is_d2h: bool) -> None: mesh = getattr(getattr(self.arrays[0], "sharding", None), "mesh", None) if mesh is None: raise ValueError("Arrays must be sharded on a Mesh for FFI weight sync.") + mesh = _reduce_mesh(mesh) slice_byte_sizes = [ int(np.prod(arr.sharding.shard_shape(arr.shape)) * arr.dtype.itemsize) @@ -291,8 +465,17 @@ def _init_ffi_transport(self, *, is_d2h: bool) -> None: ) task_mesh_shape = tuple(mesh.shape[a] for a in mesh.axis_names) - global_ids = jnp.array( - [d.id for d in mesh.devices.flatten()], dtype=jnp.int32 + # Mesh POSITION, not device id. The controller indexes a source shard by + # its position in the mesh (`_get_global_indices` walks + # physical_mesh_shape), while the native layer keys staging off whatever we + # pass here -- slot = shard_idx % num_shards, submanager = shard_idx / + # num_shards, and SetGlobalShardIndices records it as the global index. + # create_device_mesh reorders devices for topology (a 2x2x2 v5p slice comes + # back as ids [0,1,3,2,6,7,5,4]), so keying off d.id labels each slice with + # the wrong global index. A 2x2x1 slice happens to be identity-ordered, + # which is why this only ever showed up multi-host. + global_ids = jnp.arange( + mesh.devices.size, dtype=jnp.int32 ).reshape(task_mesh_shape) shard_idx = jax.device_put( global_ids, @@ -302,10 +485,27 @@ def _init_ffi_transport(self, *, is_d2h: bool) -> None: ) src_devices = mesh.devices.flatten() - num_processes = len( - set(getattr(d, "process_index", 0) for d in src_devices) + devices_per_host_env = os.environ.get("RAIDEN_DEVICES_PER_HOST") + if devices_per_host_env: + devices_per_host = int(devices_per_host_env) + elif self._is_proxy: + # In Pathways, process_index is always 0 for proxy devices. Default to + # 4 devices/host for standard Cloud TPU VM topologies. + devices_per_host = min(4, len(src_devices)) + else: + devices_per_host = _devices_per_host(list(src_devices)) + # Loud on purpose: a wrong value here is silent, and costs exactly the + # shards of every host but one. + logging.warning( + "raiden ffi: %d device(s), devices_per_host=%d (task_id=%s" + " process_index=%s local_hardware_id=%s)", + len(src_devices), + devices_per_host, + sorted({getattr(d, "task_id", None) for d in src_devices}), + sorted({getattr(d, "process_index", None) for d in src_devices}), + sorted({getattr(d, "local_hardware_id", None) for d in src_devices}, + key=str), ) - devices_per_host = len(src_devices) // max(1, num_processes) if is_d2h: logging.info( @@ -314,7 +514,7 @@ def _init_ffi_transport(self, *, is_d2h: bool) -> None: len(self.arrays), devices_per_host, ) - ws_info = _raiden_ffi.init_weight_synchronizer_and_d2h( + ws_info = raiden_ffi.init_weight_synchronizer_and_d2h( device_arrays=self.arrays, shard_idx=shard_idx, mesh=mesh, @@ -331,16 +531,52 @@ def _init_ffi_transport(self, *, is_d2h: bool) -> None: len(self.arrays), devices_per_host, ) - ws_info = _raiden_ffi.init_weight_synchronizer( - device_array=self.arrays[0], - shard_idx=shard_idx, - mesh=mesh, - slice_byte_sizes=slice_byte_sizes_sharded, - parallelism=self._parallelism, - num_layers=len(self.arrays), - listener_port=0, - num_shards=devices_per_host, + + # TODO(b/557061810): Re-enable this once the bug is fixed and the FFI + # call is verified to work. + # ws_info = _raiden_ffi.init_weight_synchronizer( + # device_array=self.arrays[0], + # shard_idx=shard_idx, + # mesh=mesh, + # slice_byte_sizes=slice_byte_sizes_sharded, + # parallelism=self._parallelism, + # num_layers=len(self.arrays), + # listener_port=0, + # num_shards=devices_per_host, + # ) + + @compute_on.compute_on( + compute_type="device_host", + out_memory_spaces=jax.memory.Space.Device, ) + def _local_init(anchor, s_idx, sizes): + axis_names = mesh.axis_names + out_shape = tuple([1] * len(axis_names)) + (6,) + return jax.ffi.ffi_call( + "init_weight_synchronizer", + jax.ShapeDtypeStruct(out_shape, jnp.int32), + has_side_effect=True, + )( + anchor, + s_idx, + sizes, + local_port=np.int32(0), + parallelism=np.int32(self._parallelism), + num_layers=np.int32(len(self.arrays)), + listener_port=np.int32(0), + num_shards=np.int32(devices_per_host), + ) + + ws_info = jax.shard_map( + _local_init, + mesh=mesh, + in_specs=( + self.arrays[0].sharding.spec, + jax.sharding.PartitionSpec(*mesh.axis_names), + jax.sharding.PartitionSpec(None), + ), + out_specs=jax.sharding.PartitionSpec(*mesh.axis_names, None), + )(self.arrays[0], shard_idx, slice_byte_sizes_sharded) local_ws_info = multihost_utils.global_array_to_host_local_array( ws_info, @@ -356,26 +592,39 @@ def _init_ffi_transport(self, *, is_d2h: bool) -> None: ip = unpack_ip(row) self._ips.append(f"{ip}:{int(row[4])}") listeners.append(f"{ip}:{int(row[5])}") + self._listeners = listeners self._unique_listeners = [] for listener in listeners: if listener not in self._unique_listeners: self._unique_listeners.append(listener) + # The controller addresses source shard j by shards[j], indexing by MESH + # position; this list is assembled by process_allgather, which orders by + # PROCESS. Under Pathways there is a single client process, so verify both + # that all devices reported and that entry j lines up with mesh device j. + logging.warning( + "raiden ffi endpoints: %d row(s) for %d mesh device(s); mesh ids=%s;" + " endpoints=%s", + len(gathered_ws_info), + len(src_devices), + [d.id for d in src_devices], + self._ips, + ) self._ffi_mesh = mesh self._ffi_shard_idx = shard_idx def _ffi_h2d(self) -> None: - if _raiden_ffi is None: + raiden_ffi = _get_raiden_ffi() + if raiden_ffi is None: raise RuntimeError( "weight_synchronizer_ffi is not available for FFI weight sync." ) if self._ffi_mesh is None or self._ffi_shard_idx is None: raise RuntimeError(f"{self.job_name}: bind() must run before h2d()") self.arrays = list( - _raiden_ffi.multi_h2d(self.arrays, self._ffi_shard_idx, self._ffi_mesh) + raiden_ffi.multi_h2d(self.arrays, self._ffi_shard_idx, self._ffi_mesh) ) - for arr in self.arrays: - arr.block_until_ready() + jax.block_until_ready(self.arrays) def bind(self, state: Any) -> None: """Binds this host's weights, or rebinds them after a training step.""" @@ -384,8 +633,14 @@ def bind(self, state: Any) -> None: # copies in host memory during rebinds. self.names = [] self.arrays = [] + if self._host_stage: + state = to_host_cpu_state(state) + _log_rss("bind:after_host_stage") self.names, self.arrays = _filter_bindable( - *flatten_weights(state), allow_proxy=self._is_proxy + # Proxy arrays are bindable only under FFI, which binds them in place. + # Host staging turns them into CPU arrays, and a non-Pathways process + # should not be seeing them at all. + *flatten_weights(state), allow_proxy=self._use_ffi ) del state _log_rss("bind:after_flatten") @@ -395,7 +650,7 @@ def bind(self, state: Any) -> None: len(self.arrays), self._is_proxy, ) - if self._is_proxy: + if self._use_ffi: self._ips = [] self._unique_listeners = [] if self._auto_h2d: @@ -407,7 +662,8 @@ def bind(self, state: Any) -> None: self._unique_listeners, ) return - if _ws_lib is None: + ws_lib = _get_ws_lib() + if ws_lib is None: return if self._sync is None: logging.info( @@ -415,7 +671,7 @@ def bind(self, state: Any) -> None: self.job_name, len(self.arrays), ) - self._sync = _ws_lib.WeightSynchronizer( + self._sync = ws_lib.WeightSynchronizer( self.arrays, local_port=0, parallelism=self._parallelism, @@ -441,12 +697,12 @@ def bind(self, state: Any) -> None: _log_rss("bind:after_native_rebind") def _require_sync(self, op: str) -> Any: - if self._sync is None and not self._is_proxy: + if self._sync is None and not self._use_ffi: raise RuntimeError(f"{self.job_name}: bind() must run before {op}") return self._sync def d2h(self) -> None: - if self._is_proxy: + if self._use_ffi: try: self._init_ffi_transport(is_d2h=True) logging.info( @@ -468,13 +724,119 @@ def d2h(self) -> None: def h2d(self) -> None: if not self.bound: raise RuntimeError(f"{self.job_name}: bind() must run before h2d()") - if self._is_proxy: + if self._use_ffi: self._ffi_h2d() return if self._sync is not None: self._sync.h2d() jax.block_until_ready(self.arrays) + def apply_to_runner(self, runner: Any) -> None: + """Applies updated arrays after H2D to the runner's state_leaves and state.""" + if runner is None or not self.arrays: + return + if not hasattr(runner, "state_leaves") or runner.state_leaves is None: + raise ValueError( + f"{self.job_name}: runner does not have a valid 'state_leaves' attribute." + ) + if not hasattr(runner, "state") or runner.state is None: + raise ValueError( + f"{self.job_name}: runner does not have a valid 'state' attribute." + ) + + new_leaves = list(runner.state_leaves) + runner_leaves_with_path = list( + jax.tree_util.tree_leaves_with_path(runner.state) + ) + if len(new_leaves) != len(runner_leaves_with_path): + raise RuntimeError( + f"{self.job_name}: runner.state_leaves length ({len(new_leaves)}) " + f"does not match runner.state leaves count ({len(runner_leaves_with_path)})." + ) + + name_to_entry = {} + for idx, (name, arr) in enumerate(zip(self.names, self.arrays)): + norm = _normalize_param_name(name) + name_to_entry[norm] = (idx, name, arr) + canon = _canonicalize_param_name(name) + if canon: + name_to_entry[canon] = (idx, name, arr) + + matched_indices = set() + for i, (path, leaf) in enumerate(runner_leaves_with_path): + p_str = jax.tree_util.keystr(path) + norm_p = _normalize_param_name(p_str) + canon_p = _canonicalize_param_name(p_str) + + entry = None + if norm_p in name_to_entry: + entry = name_to_entry[norm_p] + elif canon_p in name_to_entry: + entry = name_to_entry[canon_p] + else: + for k, v in name_to_entry.items(): + if norm_p.endswith(k) or (canon_p and canon_p.endswith(k)): + entry = v + break + + if entry is not None: + idx, orig_name, arr = entry + leaf_arr = getattr(leaf, "value", leaf) + if hasattr(leaf_arr, "shape") and leaf_arr.shape != arr.shape: + raise ValueError( + f"Shape mismatch for parameter '{orig_name}' (runner path '{p_str}'): " + f"runner shape {leaf_arr.shape} vs synchronizer shape {arr.shape}" + ) + new_leaves[i] = arr + matched_indices.add(idx) + + if len(matched_indices) != len(self.arrays): + unmatched = [ + self.names[j] + for j in range(len(self.arrays)) + if j not in matched_indices + ] + raise RuntimeError( + f"{self.job_name}: Not all synchronizer arrays were matched in runner.state! " + f"Matched {len(matched_indices)} of {len(self.arrays)} arrays. " + f"Unmatched {len(unmatched)} parameters, e.g.: {unmatched[:10]}" + ) + + runner.state_leaves = tuple(new_leaves) + runner.state = jax.tree_util.tree_unflatten( + jax.tree_util.tree_structure(runner.state), new_leaves + ) + logging.info( + "%s apply_to_runner: successfully applied %d arrays to runner state and state_leaves (total runner leaves: %d).", + self.job_name, + len(matched_indices), + len(new_leaves), + ) + + def release_host_arrays(self) -> None: + """Drops the staged host copy between rounds. + + Called by the pinned MaxText (maxtext_engine.py:1129), not by anything in + this repo -- grep site-packages before deleting. Host-staged path only; a + no-op under FFI, which binds device arrays in place. Clears `names` + alongside `arrays` so `bound`/`active` and every zip(names, arrays) + consumer stay consistent; the next round rebinds. + """ + if not self._host_stage: + return + self.names = [] + self.arrays = [] + gc.collect() + + def work_unit_metadata_all(self) -> List[weight_sync.WorkUnitMetadata]: + """Returns work unit metadata for registration. + + In proxy/FFI mode, control_addr contains all comma-separated unique listener + addresses. The coordinator registers a single work unit and the Raiden + controller broadcasts to all listeners in parallel. + """ + return [self.work_unit_metadata()] + def metrics(self) -> dict: return self._sync.get_metrics() if self._sync else {} @@ -489,27 +851,34 @@ def total(arr): for name, arr in list(zip(self.names, self.arrays))[:sample] } head["__grand_total__"] = float(sum(total(a) for a in self.arrays)) + # Registration pairs tensors by position, so the totals only compare when + # both sides bound the same set. Check these before trusting a mismatch. + head["__tensor_count__"] = len(self.arrays) + head["__element_count__"] = int(sum(a.size for a in self.arrays)) return head def work_unit_metadata(self) -> weight_sync.WorkUnitMetadata: - variables = tuple( - _tensor_metadata(name, arr, idx) - for idx, (name, arr) in enumerate(zip(self.names, self.arrays)) - ) - mesh_axes: tuple = () - mesh_shape = None + mesh = None for arr in self.arrays: mesh = getattr(getattr(arr, "sharding", None), "mesh", None) if mesh is not None: - mesh_axes = tuple(mesh.axis_names) - mesh_shape = tuple(mesh.shape[a] for a in mesh.axis_names) break - if mesh_shape is None: - mesh_axes = ("fsdp",) - mesh_shape = (1,) - if self._is_proxy: + if mesh is None: + mesh_axes, mesh_shape = ("fsdp",), (1,) + else: + # Advertise the same mesh the shards were built on; see _reduce_mesh. + mesh = _reduce_mesh(mesh) + mesh_axes = tuple(mesh.axis_names) + mesh_shape = tuple(int(mesh.shape[a]) for a in mesh.axis_names) + variables = tuple( + _tensor_metadata(name, arr, idx) + for idx, (name, arr) in enumerate(zip(self.names, self.arrays)) + ) + if self._use_ffi or self._is_proxy: shards = tuple(self._ips) - control_addr = self._unique_listeners[0] if self._unique_listeners else "" + control_addr = ( + ",".join(self._unique_listeners) if self._unique_listeners else "" + ) else: data_addr = f"{self.ip}:{self._sync.local_port}" if self._sync else "" control_addr = ( @@ -534,3 +903,29 @@ def work_unit_metadata(self) -> weight_sync.WorkUnitMetadata: variables=variables, mesh_axes=mesh_axes or None, ) + + +def patch_raiden_worker_sync() -> None: + """Monkey-patches tpu_inference.rl.raiden_worker_sync.RaidenWorkerSync to delegate apply_to_runner.""" + try: + import tpu_inference.rl.raiden_worker_sync as rws + if getattr(rws.RaidenWorkerSync, "_patched_by_tunix", False): + return + orig_apply = getattr(rws.RaidenWorkerSync, "apply_to_runner", None) + + def _patched_apply_to_runner(self, runner: Any) -> None: + if self._sync is not None and hasattr(self._sync, "apply_to_runner"): + self._sync.apply_to_runner(runner) + return + if orig_apply is not None: + orig_apply(self, runner) + + rws.RaidenWorkerSync.apply_to_runner = _patched_apply_to_runner + rws.RaidenWorkerSync._patched_by_tunix = True + logging.info("Successfully patched RaidenWorkerSync.apply_to_runner with Tunix delegation.") + except (ImportError, AttributeError) as e: + logging.debug("tpu_inference not available to patch: %s", e) + + +# Patch upon import so workers have delegation enabled automatically +patch_raiden_worker_sync() diff --git a/tunix/experimental/weight_sync/raiden_weight_sync_delegate.py b/tunix/experimental/weight_sync/raiden_weight_sync_delegate.py index 61546f2d1..81eb1762b 100644 --- a/tunix/experimental/weight_sync/raiden_weight_sync_delegate.py +++ b/tunix/experimental/weight_sync/raiden_weight_sync_delegate.py @@ -16,11 +16,13 @@ from __future__ import annotations +import asyncio import os -from typing import Any, List +from typing import Any, List, Mapping, Optional from absl import logging from tunix.experimental.weight_sync import raiden_synchronizer +from tunix.experimental.weight_sync import weight_sync_coordinator class RaidenWeightSyncDelegate: @@ -36,14 +38,33 @@ class RaidenWeightSyncDelegate: abort after a partial weight_sync cannot restore the previous weights. """ - def __init__(self, *args, worker_index: int = 0, **kwargs): + def __init__( + self, + *args, + job_name: Optional[str] = None, + server_id: Optional[str] = None, + worker_index: int = 0, + **kwargs, + ): super().__init__(*args, **kwargs) + # Raiden partitions the weights across every unit sharing a job_name, so + # replicas that all call themselves "rollout" get a slice each instead of a + # copy each. server_id is already unique per replica and shared across the + # hosts within one, which is exactly the grouping job_name needs. + self.job_name = ( + job_name or server_id or getattr(self, "server_id", None) or "rollout" + ) + self.server_id = server_id or getattr(self, "server_id", None) self._synchronizers: List[Any] = [ raiden_synchronizer.RaidenSynchronizer( - "rollout", worker_index=worker_index, auto_h2d=True + self.job_name, + worker_index=worker_index, + auto_h2d=True, ) ] self._version = 0 + self._sync_lock = asyncio.Lock() + self._tracker = weight_sync_coordinator.WorkerRoundTracker() def is_bounded( self, @@ -69,30 +90,71 @@ async def get_weight_sync_metadata(self, **kwargs) -> Any: del kwargs return [s.work_unit_metadata() for s in self._synchronizers] + def _has_round(self, sync_request: Any) -> bool: + extra = getattr(sync_request, "extra_config", None) or {} + return extra.get("req_id") is not None + async def pre_weight_sync(self, sync_request: Any = None, **kwargs) -> Any: """Pre-sync phase hook executed before weight transfer begins.""" - del sync_request, kwargs - return True + del kwargs + async with self._sync_lock: + if self._has_round(sync_request): + if not self._tracker.admit(sync_request, "prepared"): + return True + self._tracker.complete(sync_request, "prepared") + return True async def weight_sync(self, sync_request: Any = None, **kwargs) -> Any: """Executes weight installation on device from host staging buffer.""" del kwargs - for sync in self._synchronizers: - if not sync.bound: - raise RuntimeError("bind_weight_sync must run before weight_sync") - # auto_h2d installs chunks as they arrive; this call is the round's - # awaited install, so completion is guaranteed before checksums/post. - sync.h2d() - if os.environ.get("VERIFY_WEIGHTS", "").lower() == "true": - logging.info("destination checksums: %s", sync.checksums()) - version = getattr(sync_request, "policy_version", 0) - self._version = version if version else self._version + 1 - return self._version + async with self._sync_lock: + if self._has_round(sync_request): + if not self._tracker.admit(sync_request, "h2d_done"): + return self._version + + for sync in self._synchronizers: + if not sync.bound: + raise RuntimeError("bind_weight_sync must run before weight_sync") + # auto_h2d installs chunks as they arrive; this call is the round's + # awaited install, so completion is guaranteed before checksums/post. + sync.h2d() + if os.environ.get("VERIFY_WEIGHTS", "").lower() == "true": + logging.info("destination checksums: %s", sync.checksums()) + version = getattr(sync_request, "policy_version", 0) + self._version = version if version else self._version + 1 + + if self._has_round(sync_request): + self._tracker.complete(sync_request, "h2d_done") + + return self._version async def post_weight_sync(self, sync_request: Any = None, **kwargs) -> Any: """Post-sync phase hook executed after weight installation completes.""" - del sync_request, kwargs - if os.environ.get("VERIFY_WEIGHTS", "").lower() == "true": - for sync in self._synchronizers: - logging.info("raiden metrics: %s", sync.metrics()) - return True + del kwargs + async with self._sync_lock: + if self._has_round(sync_request): + if not self._tracker.admit(sync_request, "committed"): + return True + + if os.environ.get("VERIFY_WEIGHTS", "").lower() == "true": + for sync in self._synchronizers: + logging.info("raiden metrics: %s", sync.metrics()) + + if self._has_round(sync_request): + self._tracker.complete(sync_request, "committed") + + return True + + async def abort_weight_sync(self, sync_request: Any = None, **kwargs) -> Any: + """Safely handles abort of weight sync round.""" + del kwargs + async with self._sync_lock: + if self._has_round(sync_request): + if not self._tracker.admit(sync_request, "aborted"): + return False + self._tracker.complete(sync_request, "aborted") + return True + + def get_weight_sync_status(self) -> Mapping[str, Any]: + """Reports worker-side round status for coordinator recovery checks.""" + return self._tracker.report() diff --git a/tunix/experimental/weight_sync/weight_sync.py b/tunix/experimental/weight_sync/weight_sync.py index d3d178bf1..840b79f91 100644 --- a/tunix/experimental/weight_sync/weight_sync.py +++ b/tunix/experimental/weight_sync/weight_sync.py @@ -39,6 +39,9 @@ class WeightSyncMode(str, enum.Enum): RAIDEN = "raiden" +DEFAULT_WEIGHT_SYNC_MODE = WeightSyncMode.FALLBACK + + @dataclasses.dataclass(frozen=True) class WorkUnitId: """Transport-neutral identity for one participant's data work unit. @@ -76,9 +79,13 @@ class TensorMetadata: layout: Layout mapping. item_size: Bytes per element. layer_idx: Stable batching ordinal. - sharding_spec: One mesh axis name per TENSOR dimension, empty string where - that dimension is replicated. This is the subset of JAX `PartitionSpec` - used by the Tunix/JAX adapters: `P(None, "y")` is `("", "y")`. Together + sharding_spec: The mesh axis name sharding each TENSOR dimension, empty + string where that dimension is replicated. This is the subset of JAX + `PartitionSpec` used by the Tunix/JAX adapters: `P(None, "y")` is + `("", "y")`. A dimension sharded over the product of several axes -- JAX + `P(("x", "y"))`, as MoE weights get when tensor and attention-data + parallelism are combined -- is the axes joined by commas, major first: + `("x,y",)`. Together with the work unit's physical `mesh_axes`, it maps device coordinates onto the variable's logical mesh. A concrete transport must reject forms its wire representation cannot encode. @@ -129,7 +136,8 @@ def __post_init__(self) -> None: f"variable {self.name!r}: sharding_spec {self.sharding_spec} must" f" have rank {rank}" ) - named_axes = [axis for axis in self.sharding_spec if axis] + named_axes = [a for axis in self.sharding_spec for a in axis.split(",") + if a] if len(named_axes) != len(set(named_axes)): raise ValueError( f"variable {self.name!r}: a mesh axis may not shard two tensor" diff --git a/tunix/experimental/weight_sync/weight_sync_coordinator.py b/tunix/experimental/weight_sync/weight_sync_coordinator.py index 2532a93a6..2ca82dc3a 100644 --- a/tunix/experimental/weight_sync/weight_sync_coordinator.py +++ b/tunix/experimental/weight_sync/weight_sync_coordinator.py @@ -985,6 +985,7 @@ async def record_workers(final_error: str = "") -> None: except asyncio.CancelledError: raise except Exception as e: # pylint: disable=broad-except + logging.error("pre-quiesce setup failed: %s", e, exc_info=True) failures.append(f"pre-quiesce setup: {e!r}") raise fail( "bind/metadata/source-prepare failed before any destination was" @@ -1009,6 +1010,34 @@ async def record_workers(final_error: str = "") -> None: raise fail("metadata collection returned an empty side") source_units = tuple(m.unit for m in src_metadata) destination_units = tuple(m.unit for m in dst_metadata) + # Log the identities: Raiden partitions the weights across units sharing + # a job_name and broadcasts across distinct ones, so this is what decides + # whether a replica gets a copy or a slice. warning, not info -- absl + # drops INFO at its default verbosity and nothing else records the split. + for side, metas in (("src", src_metadata), ("dst", dst_metadata)): + for m in metas: + logging.warning( + "%s unit job_name=%r job_replica_id=%r shards=%d %s", + side, + m.unit.job_name, + m.unit.job_replica_id, + len(m.shards), + list(m.shards), + ) + src_shards = sum(len(m.shards) for m in src_metadata) + for m in dst_metadata: + if src_shards and len(m.shards) != src_shards: + logging.warning( + "destination %r has %d shard(s) against the source's %d. Raiden" + " intersects the two global index spaces, so an unequal pair can" + " transfer only the overlap -- a green round that delivers part" + " of the model, with every tensor that did arrive checksumming" + " correctly. Compare __grand_total__ on both sides before" + " trusting this round.", + m.unit.job_name, + len(m.shards), + src_shards, + ) # Manifest preflight, before registration and before any downtime: # the controller pairs variables by exact name and silently skips @@ -1017,9 +1046,32 @@ async def record_workers(final_error: str = "") -> None: preflight_problems = _manifest_mismatches(src_metadata, dst_metadata) if preflight_problems: failures.extend(preflight_problems) + src_names = [] + for m in src_metadata: + src_names.extend(v.name for v in m.variables) + src_names.sort() + + dst_names = [] + for m in dst_metadata: + dst_names.extend(v.name for v in m.variables) + dst_names.sort() + logging.error( + "manifest preflight failed: %d source var(s), %d destination" + " var(s), %d problem(s)\n" + " source sample:\n %s\n" + " destination sample:\n %s\n" + " problems (first 20):\n %s", + len(src_names), + len(dst_names), + len(preflight_problems), + "\n ".join(src_names[:8]), + "\n ".join(dst_names[:8]), + "\n ".join(preflight_problems[:20]), + ) raise fail( "manifest preflight failed before any destination was quiesced;" - " no rollback needed" + f" no rollback needed ({len(preflight_problems)} problems, first:" + f" {preflight_problems[0]})" ) loop = asyncio.get_running_loop() @@ -1159,7 +1211,7 @@ async def record_workers(final_error: str = "") -> None: except asyncio.CancelledError: raise except Exception as e: # pylint: disable=broad-except - # The call returned (by raising): the thread is done, rollback is safe. + logging.error("transfer raised exception: %s", e, exc_info=True) transfer_in_flight = False failures.append(f"transfer: {e!r}") state = await self._rollback(destinations, prepared_request, failures) diff --git a/tunix/experimental/worker/rollout_worker.py b/tunix/experimental/worker/rollout_worker.py index 3cd8c96cd..79b71ef7c 100644 --- a/tunix/experimental/worker/rollout_worker.py +++ b/tunix/experimental/worker/rollout_worker.py @@ -530,11 +530,10 @@ async def get_weight_sync_metadata(self, **kwargs) -> Any: async def abort_weight_sync(self, sync_request: Any = None, **kwargs) -> Any: """Discards the round and resumes serving the previous weights.""" - self.manager.resume_all() - self.manager.reopen_admission() + res = await self.manager.abort_weight_sync(sync_request, **kwargs) self.state = WorkerState.READY self._record_round(sync_request, "aborted") - return None + return res async def get_weight_sync_status(self, **kwargs) -> Any: """Returns this worker's view of the current weight sync round.""" diff --git a/tunix/generate/utils.py b/tunix/generate/utils.py index 98806be74..fb4f6a7fe 100644 --- a/tunix/generate/utils.py +++ b/tunix/generate/utils.py @@ -1864,7 +1864,18 @@ def intersect_trees( source=final_source, target=traverse_util.unflatten_dict(dst_shardings_flat), ) - nnx.update(dst_state, resharded_weights) + if hasattr(dst_state, "flat_state"): + flat_resharded = traverse_util.flatten_dict(resharded_weights) + for path, var in dst_state.flat_state(): + key_tuple = tuple(path) + if key_tuple in flat_resharded: + new_val = flat_resharded[key_tuple] + if hasattr(var, "value"): + var.value = getattr(new_val, "value", new_val) + else: + var[...] = getattr(new_val, "value", new_val) + else: + nnx.update(dst_state, resharded_weights) def resolve_parallelism_sizes( diff --git a/tunix/utils/maxtext_utils.py b/tunix/utils/maxtext_utils.py index 33abb356d..26e519f3c 100644 --- a/tunix/utils/maxtext_utils.py +++ b/tunix/utils/maxtext_utils.py @@ -59,7 +59,10 @@ def build_maxtext_config( warmup_steps_fraction: float = 0.0, load_parameters_path: str = "", padded_moe_mlp_dim: int = 0, + base_num_kv_heads: int = 0, base_output_directory: str = "", + rollout_mesh_tp: int = 0, + prefuse_moe_weights: bool | None = None, ) -> Any: """Builds the MaxText HyperParameters the training engine runs on.""" pyconfig, _, _ = maxtext_modules() @@ -86,13 +89,35 @@ def build_maxtext_config( f"run_name={worker_id or 'tunix_maxtext'}", f"base_output_directory={output_dir}", f"enable_checkpointing={enable_checkpointing}", + "skip_jax_distributed_system=True", ] if load_parameters_path: argv.append(f"load_parameters_path={load_parameters_path}") + + if rollout_mesh_tp <= 0: + rollout_mesh_tp = int( + os.environ.get("ROLLOUT_TENSOR_PARALLEL_SIZE", 0) + or os.environ.get("ROLLOUT_MESH_TP", 0) + or 0 + ) + + if not padded_moe_mlp_dim and rollout_mesh_tp > 0: + try: + from maxtext.integration.vllm.moe_padding import compute_padded_moe_mlp_dim + tmp_cfg = pyconfig.initialize(argv) + base_dim = getattr(tmp_cfg, "base_moe_mlp_dim", None) or getattr(tmp_cfg, "moe_intermediate_size", None) + if base_dim: + padded_moe_mlp_dim = compute_padded_moe_mlp_dim(base_dim, rollout_mesh_tp) + logging.info("Auto-computed padded_base_moe_mlp_dim=%d for rollout_mesh_tp=%d", padded_moe_mlp_dim, rollout_mesh_tp) + except Exception as e: + logging.warning("Could not auto-compute padded_base_moe_mlp_dim: %s", e) + + if prefuse_moe_weights is None: + prefuse_moe_weights = os.environ.get("PREFUSE_MOE_WEIGHTS", "false").lower() in ("1", "true", "yes") + argv.extend([ "scan_layers=True", "convert_checkpoint_if_possible=False", - "skip_jax_distributed_system=True", f"per_device_batch_size={per_device_batch_size}", "gradient_accumulation_steps=1", f"max_target_length={max_prompt_length + max_response_length}", @@ -105,6 +130,12 @@ def build_maxtext_config( if padded_moe_mlp_dim else [] ), + # The vLLM rollout replicates KV heads up to kv_tp_size (tp*ep) when the + # model has fewer -- see maxtext_vllm_adapter. Weight sync pairs by name, + # so the trainer must build the same shape. Prefer attention DP on the + # rollout instead, which avoids the replication entirely; this is the + # fallback when that is not available. + *([f"base_num_kv_heads={base_num_kv_heads}"] if base_num_kv_heads else []), f"ici_tensor_parallelism={mesh_tp}", f"ici_expert_parallelism={mesh_expert}", f"learning_rate={learning_rate}", @@ -115,6 +146,7 @@ def build_maxtext_config( "enable_tensorboard=False", "record_internal_nn_metrics=False", "init_weights_seed=42", + f"prefuse_moe_weights={prefuse_moe_weights}", ]) logging.info("MaxText config argv: %s", argv) return pyconfig.initialize(argv)