Internal change

init2winit/model_lib/mdlm_rope_nanodo.py

@@ -54,6 +54,7 @@ class TimestepEmbedding(nn.Module):
 
   D: int
   dtype: jnp.dtype = jnp.float32
+  param_dtype: jnp.dtype = jnp.float32
 
   @nn.compact
   def __call__(self, t_B: jax.Array):
@@ -64,9 +65,9 @@ def __call__(self, t_B: jax.Array):
     angles = t_B[:, None] * freq[None, :]
     sincos = jnp.concatenate([jnp.sin(angles), jnp.cos(angles)], axis=-1)
 
-    h = nn.Dense(self.D, dtype=self.dtype)(sincos)
+    h = nn.Dense(self.D, dtype=self.dtype, param_dtype=self.param_dtype)(sincos)
     h = nn.gelu(h)
-    h = nn.Dense(self.D, dtype=self.dtype)(h)
+    h = nn.Dense(self.D, dtype=self.dtype, param_dtype=self.param_dtype)(h)
     return h
 
 
@@ -81,23 +82,37 @@ def setup(self):
         num_embeddings=cfg.V + 1,
         features=cfg.D,
         embedding_init=cfg.embed_init,
+        dtype=cfg.dtype,
+        param_dtype=cfg.param_dtype,
     )
 
-    self.time_embed = TimestepEmbedding(D=cfg.D, dtype=cfg.dtype)
+    self.time_embed = TimestepEmbedding(
+        D=cfg.D, dtype=cfg.dtype, param_dtype=cfg.param_dtype
+    )
 
     self.blocks = [rope_nanodo.TBlock(cfg) for _ in range(cfg.N)]
     if cfg.normalization == 'layernorm':
-      self.out_ln = nn.LayerNorm(dtype=cfg.dtype, use_bias=False)
+      self.out_ln = nn.LayerNorm(
+          dtype=cfg.dtype, param_dtype=cfg.param_dtype, use_bias=False
+      )
     elif cfg.normalization == 'rmsnorm':
-      self.out_ln = nn.RMSNorm(dtype=cfg.dtype, epsilon=cfg.rmsnorm_epsilon)
+      self.out_ln = nn.RMSNorm(
+          dtype=cfg.dtype,
+          param_dtype=cfg.param_dtype,
+          epsilon=cfg.rmsnorm_epsilon,
+      )
     else:
       raise ValueError(f'Unknown normalization: {cfg.normalization}')
 
     if cfg.tie_embeddings:
       self.output_proj = None
     else:
       self.output_proj = nn.Dense(
-          cfg.V, kernel_init=cfg.embed_init, dtype=cfg.dtype, name='output_proj'
+          cfg.V,
+          kernel_init=cfg.embed_init,
+          dtype=cfg.dtype,
+          param_dtype=cfg.param_dtype,
+          name='output_proj',
       )
 
   def __call__(self, z_BxL: jax.Array, t_B: jax.Array, train: bool):
@@ -117,7 +132,7 @@ def __call__(self, z_BxL: jax.Array, t_B: jax.Array, train: bool):
       logits_BxLxV = self.output_proj(z_BxLxD)
     else:
       embed_matrix = self.embed.embedding[: cfg.V]
-      logits_BxLxV = z_BxLxD.astype(jnp.float32) @ embed_matrix.T
+      logits_BxLxV = z_BxLxD @ embed_matrix.astype(cfg.dtype).T
 
     return logits_BxLxV
 
@@ -134,6 +149,7 @@ def build_flax_module(self):
         L=self.hps['input_shape'][0],
         F=self.hps['mlp_dim'],
         dtype=utils.dtype_from_str(self.hps['computation_dtype']),
+        param_dtype=utils.dtype_from_str(self.hps['model_dtype']),
         mlp_activation=self.hps['mlp_activation'],
         normalization=self.hps['normalization'],
         qk_norm=self.hps['qk_norm'],

init2winit/model_lib/rope_nanodo.py

@@ -66,7 +66,8 @@ class DoConfig:
   embed_init: nn.initializers.Initializer = nn.initializers.variance_scaling(
       1.0, 'fan_in', 'normal', out_axis=0
   )
-  dtype: jnp.dtype = jnp.float32
+  dtype: jnp.dtype = jnp.bfloat16
+  param_dtype: jnp.dtype = jnp.float32
   rmsnorm_epsilon: float = 1e-6
   multiple_of: int = 256
   tie_embeddings: bool = True  # Whether to tie input and output embeddings
@@ -87,7 +88,11 @@ def __call__(self, x_BxLxD: jax.Array):
     cfg = self.cfg
     # Use Xavier uniform initialization explicitly
     linear = partial(
-        nn.Dense, kernel_init=cfg.kernel_init, use_bias=False, dtype=cfg.dtype
+        nn.Dense,
+        kernel_init=cfg.kernel_init,
+        use_bias=False,
+        dtype=cfg.dtype,
+        param_dtype=cfg.param_dtype,
     )
     if cfg.mlp_activation == 'gelu':
       mlp_activation = nn.gelu
@@ -176,6 +181,7 @@ def setup(self):
         kernel_init=cfg.kernel_init,
         use_bias=False,
         dtype=cfg.dtype,
+        param_dtype=cfg.param_dtype,
     )
 
     self.multilinear_query = self.multilinear(name='query')
@@ -188,13 +194,14 @@ def setup(self):
         kernel_init=cfg.kernel_init,
         use_bias=False,
         dtype=cfg.dtype,
+        param_dtype=cfg.param_dtype,
     )
     if cfg.qk_norm:
       self.eps = cfg.eps
-    attn_scale0 = jnp.log2(cfg.L**2 - cfg.L).astype(cfg.dtype)
+    attn_scale0 = jnp.log2(cfg.L**2 - cfg.L).astype(cfg.param_dtype)
     self.attn_scale = self.param(
         'attn_scale',
-        nn.initializers.constant(attn_scale0, dtype=cfg.dtype),
+        nn.initializers.constant(attn_scale0, dtype=cfg.param_dtype),
         (),
     )
 
@@ -212,7 +219,7 @@ def __call__(self, x_BxLxD: jax.Array):
       k_BxLxHxDh /= (
           jnp.linalg.norm(k_BxLxHxDh, axis=-1, keepdims=True) + self.eps
       )
-    q_BxLxHxDh = q_BxLxHxDh * self.attn_scale
+    q_BxLxHxDh = q_BxLxHxDh * self.attn_scale.astype(cfg.dtype)
 
     q_BxLxHxDh, k_BxLxHxDh = apply_rope(q_BxLxHxDh, k_BxLxHxDh, self.freqs_cis)
 
@@ -241,11 +248,15 @@ def __call__(self, in_BxLxD: jax.Array):
 
     # "pre-layernorm"
     if cfg.normalization == 'layernorm':
-      x_BxLxD = nn.LayerNorm(dtype=cfg.dtype, use_bias=False)(in_BxLxD)
+      x_BxLxD = nn.LayerNorm(
+          dtype=cfg.dtype, param_dtype=cfg.param_dtype, use_bias=False
+      )(in_BxLxD)
     elif cfg.normalization == 'rmsnorm':
-      x_BxLxD = nn.RMSNorm(dtype=cfg.dtype, epsilon=cfg.rmsnorm_epsilon)(
-          in_BxLxD
-      )
+      x_BxLxD = nn.RMSNorm(
+          dtype=cfg.dtype,
+          param_dtype=cfg.param_dtype,
+          epsilon=cfg.rmsnorm_epsilon,
+      )(in_BxLxD)
     else:
       raise ValueError(f'Unknown normalization: {cfg.normalization}')
 
@@ -268,18 +279,28 @@ def setup(self):
         num_embeddings=cfg.V,
         features=cfg.D,
         embedding_init=cfg.embed_init,
+        dtype=cfg.dtype,
+        param_dtype=cfg.param_dtype,
     )
     self.pos_embed = nn.Embed(
         num_embeddings=cfg.L,
         features=cfg.D,
         embedding_init=cfg.embed_init,
+        dtype=cfg.dtype,
+        param_dtype=cfg.param_dtype,
     )
 
     self.blocks = [TBlock(cfg) for _ in range(cfg.N)]
     if cfg.normalization == 'layernorm':
-      self.out_ln = nn.LayerNorm(dtype=cfg.dtype, use_bias=False)
+      self.out_ln = nn.LayerNorm(
+          dtype=cfg.dtype, param_dtype=cfg.param_dtype, use_bias=False
+      )
     elif cfg.normalization == 'rmsnorm':
-      self.out_ln = nn.RMSNorm(dtype=cfg.dtype, epsilon=cfg.rmsnorm_epsilon)
+      self.out_ln = nn.RMSNorm(
+          dtype=cfg.dtype,
+          param_dtype=cfg.param_dtype,
+          epsilon=cfg.rmsnorm_epsilon,
+      )
     else:
       raise ValueError(f'Unknown normalization: {cfg.normalization}')
 
@@ -288,7 +309,11 @@ def setup(self):
       self.output_proj = lambda x: self.embed.attend(x.astype(jnp.float32))
     else:
       self.output_proj = nn.Dense(
-          cfg.V, kernel_init=cfg.embed_init, dtype=cfg.dtype, name='output_proj'
+          cfg.V,
+          kernel_init=cfg.embed_init,
+          dtype=cfg.dtype,
+          param_dtype=cfg.param_dtype,
+          name='output_proj',
       )
 
   def __call__(self, y_BxL: jax.Array, train: bool):
@@ -314,6 +339,7 @@ def build_flax_module(self):
         L=self.hps['input_shape'][0],
         F=self.hps['mlp_dim'],
         dtype=utils.dtype_from_str(self.hps['computation_dtype']),
+        param_dtype=utils.dtype_from_str(self.hps['model_dtype']),
         mlp_activation=self.hps['mlp_activation'],
         normalization=self.hps['normalization'],
         qk_norm=self.hps['qk_norm'],

init2winit/model_lib/test_models.py

@@ -25,6 +25,7 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import flax.linen as nn
+from init2winit import utils
 from init2winit.init_lib import initializers
 from init2winit.model_lib import model_utils
 from init2winit.model_lib import models
@@ -1221,6 +1222,90 @@ def test_mlperf_resnet_params_types(self):
         }
     self.assertEqual(model.params_types, expected)
 
+  @parameterized.product(
+      model_name=['mdlm_rope_nanodo', 'rope_nanodo']
+  )
+  def test_nanodo_parameter_and_activation_dtypes(self, model_name):
+    """Verify that parameters are in model_dtype and activations are in computation_dtype."""
+    model_dtype_str = 'float32'
+    computation_dtype_str = 'bfloat16'
+    model_dtype = utils.dtype_from_str(model_dtype_str)
+    computation_dtype = utils.dtype_from_str(computation_dtype_str)
+
+    model_cls = models.get_model(model_name)
+    hps = copy.deepcopy(
+        training_algorithm.OptaxTrainingAlgorithm.get_default_training_hparams()
+    )
+    hps.update(models.get_model_hparams(model_name))
+    hps.update(DATA_HPS[model_name])
+    hps.model_dtype = model_dtype_str
+    hps.computation_dtype = computation_dtype_str
+
+    model = model_cls(
+        hps,
+        dataset_meta_data={'shift_inputs': True, 'causal': True},
+        loss_name=LOSS_NAME[model_name],
+        metrics_name=METRICS_NAME[model_name],
+    )
+
+    rng = jax.random.PRNGKey(0)
+    initializer = initializers.get_initializer('noop')
+    params, _ = model.initialize(initializer, hps, rng, metrics_logger=None)
+
+    # 1. Verify all parameters are initialized in model_dtype
+    mismatched_params = []
+    flat_params, _ = jax.tree_util.tree_flatten_with_path(params)
+    for path, val in flat_params:
+      if hasattr(val, 'dtype') and val.dtype != model_dtype:
+        mismatched_params.append((
+            jax.tree_util.keystr(path, simple=True, separator='/'),
+            str(val.dtype),
+        ))
+
+    self.assertEmpty(
+        mismatched_params,
+        msg=(
+            f'Model {model_name} parameters not in model_dtype'
+            f' {model_dtype_str}\n: {mismatched_params}'
+        ),
+    )
+
+    # 2. Run forward pass and capture intermediates to verify activation dtypes
+    fake_inputs = _get_fake_inputs_for_initialization(model, hps)
+    _, state = model.flax_module.apply(
+        {'params': params},
+        *fake_inputs,
+        train=False,
+        capture_intermediates=True,
+    )
+
+    # Verify all captured intermediates (activations) are in computation_dtype
+    # Note: This checks Embed outputs, TBlock outputs, Attention outputs,
+    # Mlp outputs, etc.
+    intermediates = state.get('intermediates', {})
+    self.assertNotEmpty(
+        intermediates, msg=f'No intermediates captured for {model_name}'
+    )
+
+    # We check all leaves in the intermediates tree and report detailed
+    # mismatches
+    mismatched = []
+    flat_intermediates, _ = jax.tree_util.tree_flatten_with_path(intermediates)
+    for path, val in flat_intermediates:
+      if hasattr(val, 'dtype') and val.dtype != computation_dtype:
+        mismatched.append((
+            jax.tree_util.keystr(path, simple=True, separator='/'),
+            str(val.dtype),
+        ))
+
+    self.assertEmpty(
+        mismatched,
+        msg=(
+            f'Model {model_name} has mismatched intermediate dtypes:\n'
+            f' {mismatched}'
+        ),
+    )
+
 
 if __name__ == '__main__':
   absltest.main()