From 952efe07175a33c554b9537d7a48ed61bcd33319 Mon Sep 17 00:00:00 2001 From: Vincent Roulet Date: Thu, 4 Jun 2026 09:21:38 -0700 Subject: [PATCH] Internal change PiperOrigin-RevId: 926726877 --- init2winit/model_lib/mdlm_rope_nanodo.py | 30 +++++++-- init2winit/model_lib/rope_nanodo.py | 50 ++++++++++---- init2winit/model_lib/test_models.py | 85 ++++++++++++++++++++++++ 3 files changed, 146 insertions(+), 19 deletions(-) diff --git a/init2winit/model_lib/mdlm_rope_nanodo.py b/init2winit/model_lib/mdlm_rope_nanodo.py index 74c3d80b..f45fdd31 100644 --- a/init2winit/model_lib/mdlm_rope_nanodo.py +++ b/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,15 +82,25 @@ 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}') @@ -97,7 +108,11 @@ def setup(self): 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'], diff --git a/init2winit/model_lib/rope_nanodo.py b/init2winit/model_lib/rope_nanodo.py index bb2d7348..ffda5592 100644 --- a/init2winit/model_lib/rope_nanodo.py +++ b/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'], diff --git a/init2winit/model_lib/test_models.py b/init2winit/model_lib/test_models.py index 0a19e153..5e5ebc3b 100644 --- a/init2winit/model_lib/test_models.py +++ b/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()