fix typos & simplify model loader

Author: Kipsora

python/sgl_jax/srt/model_loader/loader.py

@@ -7,8 +7,8 @@
 import huggingface_hub
 import jax
 import jax.numpy as jnp
-from flax import nnx
 import numpy as np
+from flax import nnx
 
 from sgl_jax.srt.configs.load_config import LoadConfig, LoadFormat
 from sgl_jax.srt.configs.model_config import ModelConfig
@@ -53,7 +53,7 @@ def init_new(cls, model_config: ModelConfig):
             )
 
     def __init__(
-        self, load_config: LoadConfig, rngs: jax.Array, mesh: jax.sharding.Mesh
+        self, load_config: LoadConfig, rngs: nnx.Rngs, mesh: jax.sharding.Mesh
     ):
         super().__init__(load_config)
         self.rng = rngs
@@ -101,7 +101,7 @@ def create_model(rng: nnx.Rngs):
             nnx.update(model, sharded_state)
             return model
 
-        with self.mesh:
+        with jax.sharding.use_mesh(self.mesh):
             model = create_model(self.rng)
 
         rng_key = self.rng.default.key.value
@@ -160,15 +160,15 @@ class JAXDummyModelLoader(BaseModelLoader):
     """Model loader that will set model weights to random values for JAX models."""
 
     def __init__(
-        self, load_config: LoadConfig, rngs: jax.Array, mesh: jax.sharding.Mesh
+        self, load_config: LoadConfig, rngs: nnx.Rngs, mesh: jax.sharding.Mesh
     ):
         super().__init__(load_config)
         if load_config.model_loader_extra_config:
             raise ValueError(
                 f"Model loader extra config is not supported for "
                 f"load format {load_config.load_format}"
             )
-        self.rng = rngs
+        self.rngs = rngs
         self.mesh = mesh
 
     def download_model(self, model_config: ModelConfig) -> None:
@@ -204,28 +204,27 @@ def _np_gen_dtype(jdtype) -> np.dtype:
                 return np.float16
             if jdtype == jnp.bfloat16:
                 return np.float32
-            return {jnp.float16: np.float16,
-                    jnp.float32: np.float32,
-                    jnp.float64: np.float64}.get(jdtype, np.float32)
-
-        def _init_leaf(x, pspec):
-            if isinstance(x, jax.Array) and jnp.issubdtype(x.dtype, jnp.floating):
-                tgt_dtype = x.dtype
+            return {
+                jnp.float16: np.float16,
+                jnp.float32: np.float32,
+                jnp.float64: np.float64,
+            }.get(jdtype, np.float32)
+
+        def _init_leaf(x):
+            if isinstance(x, nnx.Param) and jnp.issubdtype(x.value.dtype, jnp.floating):
+                tgt_dtype = x.value.dtype
                 gen_dtype = _np_gen_dtype(tgt_dtype)
-                numel = int(np.prod(x.shape))
+                numel = int(np.prod(x.value.shape))
 
                 # Per-parameter reseed (shape-agnostic stream)
                 rng = np.random.default_rng(seed)
                 flat = rng.uniform(low, high, size=(numel,)).astype(gen_dtype)
-                arr_np = flat.reshape(x.shape)
-
-                arr_jax = jnp.asarray(arr_np, dtype=tgt_dtype)
-                if pspec is not None:
-                    arr_jax = jax.lax.with_sharding_constraint(arr_jax, pspec)
-                return arr_jax
+                x.value = jnp.asarray(flat.reshape(x.shape), dtype=tgt_dtype)
             return x
 
-        new_params = jax.tree_util.tree_map(_init_leaf, params, pspecs)
+        new_params = jax.tree.map(
+            _init_leaf, params, is_leaf=lambda x: isinstance(x, nnx.Param)
+        )
 
         # Do not alter rotary embedding caches
         def _preserve_rope_caches(path, old, new):
@@ -235,15 +234,16 @@ def _preserve_rope_caches(path, old, new):
                 return old
             return new
 
-        new_params = jax.tree_util.tree_map_with_path(_preserve_rope_caches, params, new_params)
+        new_params = jax.tree.map_with_path(_preserve_rope_caches, params, new_params)
         nnx.update(model, new_params)
         print(f"new_params: {new_params}")
+
         # Print out the shape of each param in new_params
         def _print_shape(path, x):
             if isinstance(x, jax.Array):
                 print(f"Param {path}: shape={x.shape}")
-        jax.tree_util.tree_map_with_path(_print_shape, new_params)
 
+        jax.tree.map_with_path(_print_shape, new_params)
 
     def load_model(
         self,
@@ -253,24 +253,22 @@ def load_model(
         # Initialize JAX model definition on mesh
         model_class = self._initialize_model(model_config)
 
-        def create_model(rng: nnx.Rngs):
-            model = model_class(model_config.hf_config, rng, self.mesh)
+        with jax.sharding.use_mesh(self.mesh):
+            model = model_class(model_config.hf_config, self.rngs, self.mesh)
+            # self._initialize_dummy_weights(model)
             state = nnx.state(model)
             pspecs = nnx.get_partition_spec(state)
             sharded_state = jax.lax.with_sharding_constraint(state, pspecs)
             nnx.update(model, sharded_state)
-            return model
 
-        with self.mesh:
-            model = create_model(self.rng)
-            # Assign random weights deterministically
-            self._initialize_dummy_weights(model)
+        for dev in jax.local_devices():
+            print(dev, dev.memory_stats())
 
         return model
 
 
 def get_model_loader(
-    load_config: LoadConfig, rngs: jax.Array, mesh: jax.sharding.Mesh
+    load_config: LoadConfig, rngs: nnx.Rngs, mesh: jax.sharding.Mesh
 ) -> BaseModelLoader:
     """Get a model loader based on the load format."""