address Copilot comments

Author: yhtang

jax-inference-offloading/examples/download_model.py

@@ -150,6 +150,8 @@ def main():
     key = args.kaggle_key or os.getenv("KAGGLE_KEY")
     with stdout_to_stderr():
       return download_kaggle(args.model, args.flavor, username, key)
+  else:
+    raise ValueError(f"Unknown hub: {hub}")
 
 
 if __name__ == "__main__":

jax-inference-offloading/jax_inference_offloading/transport/model/nccl_grouped.py

@@ -71,7 +71,6 @@ def __init__(
   def __call__(self, named_parameters: Dict[str, jax.Array]):
     mapping_specs = self._mapping_specs
     gateway = self._gateway
-    transports = self._transports
     transport_config = self._transport_config
 
     gateway.start_weight_transfer('grouped')

jax-inference-offloading/jax_inference_offloading/transport/model/nccl_unfused.py

@@ -30,7 +30,6 @@
 from jax_inference_offloading.controller.trainer_client import TrainerClient
 from jax_inference_offloading.sharding import PolymorphicMesh
 from jax_inference_offloading.transport.tensor.nccl_star import NcclStarTransport
-from jax_inference_offloading.models.mapping_util import _proto_to_slice
 from jax_inference_offloading.timer import Timer
 
 logger = logging.getLogger(__name__)

jax-inference-offloading/jax_inference_offloading/transport/tensor/nccl_base.py

@@ -30,13 +30,6 @@
 except ImportError:
   has_torch = False
 
-try:
-  import jax
-  import jax.numpy as jnp
-  has_jax = True
-except ImportError:
-  has_jax = False
-
 
 logger = getLogger(__name__)
 

jax-inference-offloading/jax_inference_offloading/transport/tensor/nccl_star.py

@@ -31,12 +31,6 @@
 except ImportError:
   has_torch = False
 
-try:
-  import jax
-  import jax.numpy as jnp
-  has_jax = True
-except ImportError:
-  has_jax = False
 
 logger = getLogger(__name__)
 
@@ -120,6 +114,8 @@ def create_trainer_transport(cls, config: dict[str, Any]) -> List['NcclStarTrans
           comm = nccl.NcclCommunicator(world_size, unique_id, 0)  # trainer rank is at the center of the star in fan-out mode and is always rank 0
         transports.append(cls(comm=comm))
       return transports
+    else:
+      raise ValueError(f"Unknown transport mode: {config['MODE']}")
 
   @classmethod
   def create_rollout_transport(cls, config: dict[str, Any], tp_rank: int) -> 'NcclStarTransport':
@@ -141,6 +137,8 @@ def create_rollout_transport(cls, config: dict[str, Any], tp_rank: int) -> 'Nccl
       unique_id = cls.decode_nccl_id(config["UNIQUE_IDS"][star_id])
       comm = nccl.NcclCommunicator(world_size, unique_id, rank)
       return cls(comm)
+    else:
+      raise ValueError(f"Unknown transport mode: {config['MODE']}")
 
   def __init__(self, comm: nccl.NcclCommunicator):
     self._comm = comm
@@ -272,7 +270,6 @@ def scatter_grouped(self, all_buffers: List[List[Any]]) -> None:
     assert self._comm.rank_id() == 0, \
       "Star scatter must originate from the root (rank 0)."
 
-    num_peers = self._comm.size() - 1
     with cuda.Device(self._comm.device_id()):
       stream = cuda.get_current_stream().ptr
 

jax-inference-offloading/jax_inference_offloading/tunix/load_model.py

@@ -18,7 +18,6 @@
 from dataclasses import asdict, replace
 
 import jax
-from flax import nnx
 
 from tunix.models.dummy_model_creator import create_dummy_model
 

jax-inference-offloading/jax_inference_offloading/vllm/extension.py

@@ -164,12 +164,12 @@ def update_weights(self, mapping_specs: TpModelMappingSpecs):
         if sharding_specs.parallelism > 0:
           shape[sharding_specs.dim] //= sharding_specs.parallelism
 
-        # logger.warning(f'vLLM TP rank {tp_rank} receiving {param.vllm_param.name} ...')
+        logger.debug(f'vLLM TP rank {tp_rank} receiving {param.vllm_param.name} ...')
         weight = self.transport.gather(
           shape, param.vllm_param.dtype or 'bfloat16',
           sharding_specs.aux_dim, sharding_specs.aux_parallelism
         )
-        # logger.warning(f'vLLM TP rank {tp_rank} received {param.vllm_param.name} shape {weight.shape}')
+        logger.debug(f'vLLM TP rank {tp_rank} received {param.vllm_param.name} shape {weight.shape}')
         self._staged_weights.append((param.vllm_param.name, weight))
 
         # TODO: make it optional
@@ -183,8 +183,7 @@ def update_weights(self, mapping_specs: TpModelMappingSpecs):
         if sharding_specs.parallelism > 0:
           shape[sharding_specs.dim] //= sharding_specs.parallelism
 
-        raw_specs_str = ' '.join(str(sharding_specs).split('\n'))
-        logger.info(f"vLLM expecting: {param.vllm_param.name} shape {shape.tolist()} raw specs {param}")
+        logger.debug(f"vLLM expecting: {param.vllm_param.name} shape {shape.tolist()} raw specs {param}")
 
         weight = self.transport.recv(shape, param.vllm_param.dtype or 'bfloat16')
         self._staged_weights.append((param.vllm_param.name, weight))