[Trellis]Weight conversion that compatible with raiden

launch_raiden.sh

@@ -0,0 +1 @@
+/usr/local/google/home/yixuannwang/projects/tunix/tunix/experimental/examples/math_gsm8k_dist/launch_raiden.sh

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

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(

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,

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)
 

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:

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

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

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)