fix nnx.data & copy=true

Author: pathfinder-fp

python/sgl_jax/srt/layers/embeddings.py

@@ -163,7 +163,7 @@ def __init__(
                 )
             )
         else:
-            self.bias = None
+            self.bias = nnx.data(None)
 
     def tie_weights(self, embed_tokens: Embed):
         """Tie the weights with word embeddings."""

python/sgl_jax/srt/layers/layernorm.py

@@ -48,7 +48,7 @@ def __init__(
         if use_scale:
             self.scale = nnx.Param(scale_init(jax.random.PRNGKey(0), feature_shape, param_dtype))
         else:
-            self.scale = None
+            self.scale = nnx.data(None)
 
         self.num_features = num_features
         self.epsilon = epsilon

python/sgl_jax/srt/layers/linear.py

@@ -55,7 +55,7 @@ def __init__(
                 )
             )
         else:
-            self.bias = None
+            self.bias = nnx.data(None)
 
     def __call__(self, x: jax.Array) -> tuple[jax.Array, jax.Array | None]:
         """Forward pass of the linear layer."""

python/sgl_jax/srt/layers/moe.py

@@ -33,7 +33,7 @@ def __init__(
                 )
             )
         else:
-            self.bias = None
+            self.bias = nnx.data(None)
 
     def __call__(self, hidden_states: jax.Array) -> tuple[jax.Array, jax.Array | None]:
         logits = hidden_states.astype(self.weight_dtype) @ self.kernel

python/sgl_jax/srt/model_executor/model_runner.py

@@ -171,7 +171,7 @@ def jitted_run_model(
             logits_metadata,
         ):
             model_state = jax.tree_util.tree_unflatten(model_state_def, model_state_leaves)
-            model = nnx.merge(model_def, model_state)
+            model = nnx.merge(model_def, model_state, copy=True)
             return model(forward_batch, token_to_kv_pool, logits_metadata)
 
         @partial(jax.jit, static_argnames=["sampler_state_def", "mesh", "use_sort_for_toppk_minp"])
@@ -184,7 +184,7 @@ def jitted_sampler(
             *args,
         ):
             model_state = jax.tree_util.tree_unflatten(sampler_state_def, sampler_state_leaves)
-            sampler = nnx.merge(sampler_def, model_state)
+            sampler = nnx.merge(sampler_def, model_state, copy=True)
             return sampler(*args, mesh=mesh, use_sort_for_toppk_minp=use_sort_for_toppk_minp)
 
         def run_model_wrapper(forward_batch, logits_metadata):