Refactor model configuration classes to make them consistent between JAX and PyTorch, also unify initialization to be the same in both

Author: rka97

algoperf/workloads/lm/lm_jax/nanodo_model.py

@@ -12,32 +12,33 @@
 
 
 @dataclasses.dataclass
-class DoConfig:
+class ModelConfig:
   """Hyper-parameters for Transformer decoder-only."""
 
-  D: int  # model/embed dim = qkv dim
-  H: int  # num attention heads
-  L: int  # max context/sequence length
-  N: int  # number of transformer block layers
-  V: int  # vocab size
-  F: int  # FF inner dimension
+  model_dim: int  # model/embed dim = qkv dim
+  num_heads: int  # num attention heads
+  seq_len: int  # max context/sequence length
+  num_layers: int  # number of transformer block layers
+  vocab_size: int  # vocab size
+  expanded_model_dim: int  # FF inner dimension
+  multiple_of: int = 256
+  rmsnorm_epsilon: float = 1e-6
+  use_residual_scaling: bool = True
+  tie_embeddings: bool = True  # Whether to tie input and output embed
+
+  dtype: jnp.dtype = jnp.float32
   attention_init: nn.initializers.Initializer = nn.initializers.normal(stddev=0.02)
   linear_init: nn.initializers.Initializer = nn.initializers.normal(stddev=0.02)
   embed_init: nn.initializers.Initializer = nn.initializers.normal(stddev=0.02)
-  use_residual_scaling: bool = True
-  dtype: jnp.dtype = jnp.float32
-  rmsnorm_epsilon: float = 1e-6
-  multiple_of: int = 256
-  tie_embeddings: bool = True  # Whether to tie input and output embed
 
   def __post_init__(self):
-    self.residual_init = nn.initializers.normal(stddev=0.02/jnp.sqrt(2 * self.N))
+    self.residual_init = nn.initializers.normal(stddev=0.02/jnp.sqrt(2 * self.num_layers))
 
 
 class Mlp(nn.Module):
   """Multilayer perceptron with GLU activation."""
 
-  cfg: DoConfig
+  cfg: ModelConfig
 
   @nn.compact
   def __call__(self, x_BxLxD: jax.Array):
@@ -49,15 +50,15 @@ def __call__(self, x_BxLxD: jax.Array):
     #  Adjust hidden dimension to keep the number of parameters invariant to
     # the activation function used since the GLU MLP has 3 * hidden_dim * D
     # parameters instead of 2 * hidden_dim * D parameters
-    hidden_dim = cfg.F * 2 / 3
+    hidden_dim = cfg.expanded_model_dim * 2 / 3
     hidden_dim = cfg.multiple_of * (
-      (cfg.F + cfg.multiple_of - 1) // cfg.multiple_of
+      (cfg.expanded_model_dim + cfg.multiple_of - 1) // cfg.multiple_of
     )
     # Double the hidden dimension for GLU
     x_BxLx2F = linear(2 * hidden_dim)(x_BxLxD)
     # Apply GLU activation
     x_BxLxF = nn.glu(x_BxLx2F, axis=-1)
-    x_BxLxD = nn.Dense(cfg.D, use_bias=False, dtype=cfg.dtype, kernel_init=cfg.residual_init if cfg.use_residual_scaling else cfg.linear_init)(x_BxLxF)
+    x_BxLxD = nn.Dense(cfg.model_dim, use_bias=False, dtype=cfg.dtype, kernel_init=cfg.residual_init if cfg.use_residual_scaling else cfg.linear_init)(x_BxLxF)
     return x_BxLxD
 
 
@@ -109,21 +110,21 @@ def rotate_tensor(x):
 class CausalAttn(nn.Module):
   """Causal attention layer with rotary embeddings."""
 
-  cfg: DoConfig
+  cfg: ModelConfig
 
   def setup(self):
     cfg = self.cfg
-    assert cfg.D % cfg.H == 0, f'D {cfg.D} not divisible by H {cfg.H}'
-    self.Dh = cfg.D // cfg.H
+    assert cfg.model_dim % cfg.num_heads == 0, f'D {cfg.model_dim} not divisible by H {cfg.num_heads}'
+    self.Dh = cfg.model_dim // cfg.num_heads
 
     # Initialize rotary embeddings
-    self.freqs_cis = init_rope(cfg.D, cfg.L, cfg.H)
+    self.freqs_cis = init_rope(cfg.model_dim, cfg.seq_len, cfg.num_heads)
 
     # Maps D -> (H, Dh)
     self.multilinear = partial(
       nn.DenseGeneral,
       axis=-1,
-      features=(cfg.H, self.Dh),
+      features=(cfg.num_heads, self.Dh),
       kernel_init=cfg.attention_init,
       use_bias=False,
       dtype=cfg.dtype,
@@ -133,7 +134,7 @@ def setup(self):
     self.multilinear_key = self.multilinear(name='key')
     self.multilinear_value = self.multilinear(name='value')
     self.output_projection = nn.DenseGeneral(
-      features=cfg.D,
+      features=cfg.model_dim,
       name='attn_out_proj',
       # axis=(-2, -1),      #
       kernel_init=cfg.residual_init if cfg.use_residual_scaling else cfg.linear_init,
@@ -183,7 +184,7 @@ def __call__(self, x_BxLxD: jax.Array):
 class TBlock(nn.Module):
   """Transformer Block."""
 
-  docfg: DoConfig
+  docfg: ModelConfig
 
   @nn.compact
   def __call__(self, in_BxLxD: jax.Array):
@@ -208,25 +209,25 @@ def __call__(self, in_BxLxD: jax.Array):
 class TransformerDo(nn.Module):
   """Transformer decoder-only."""
 
-  docfg: DoConfig
+  docfg: ModelConfig
 
   def setup(self):
     cfg = self.docfg
     self.embed = nn.Embed(
-      num_embeddings=cfg.V,
-      features=cfg.D,
+      num_embeddings=cfg.vocab_size,
+      features=cfg.model_dim,
       embedding_init=cfg.embed_init,
     )
 
-    self.blocks = [TBlock(cfg) for _ in range(cfg.N)]
+    self.blocks = [TBlock(cfg) for _ in range(cfg.num_layers)]
     self.out_ln = nn.RMSNorm(param_dtype=cfg.dtype, epsilon=cfg.rmsnorm_epsilon)
 
     # Output projection - tied to input embeddings if configured
     if cfg.tie_embeddings:
       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.vocab_size, kernel_init=cfg.embed_init, dtype=cfg.dtype, name='output_proj'
       )
 
   def __call__(self, y_BxL: jax.Array):
@@ -255,9 +256,9 @@ def predict(self, y_BxL: jax.Array, k: int = 1):
     original_input = y_BxL
 
     # Make sure we don't exceed the model's context length
-    if seq_len + k > cfg.L:
+    if seq_len + k > cfg.seq_len:
       raise ValueError(
-        f"Total sequence length ({seq_len + k}) exceeds model's context length ({cfg.L})"
+        f"Total sequence length ({seq_len + k}) exceeds model's context length ({cfg.seq_len})"
       )
 
     # Generate k tokens autoregressively
@@ -288,25 +289,25 @@ def main():
   """Create and run the DecoderOnly Transformer model."""
   # Initialize model configuration with smaller parameters for demo
   B, L = (2, 128)  # Batch size, sequence length
-  cfg = DoConfig(D=128, H=4, L=L, N=2, V=256, F=4 * 128)
+  cfg = ModelConfig(model_dim=128, num_heads=4, seq_len=L, num_layers=2, vocab_size=256, expanded_model_dim=4 * 128)
   model = TransformerDo(cfg)
 
   # Print model info
   print('\nModel Configuration:')
-  print(f'  - Model dimension (D): {cfg.D}')
-  print(f'  - Number of heads (H): {cfg.H}')
-  print(f'  - Max sequence length (L): {cfg.L}')
-  print(f'  - Number of layers (N): {cfg.N}')
-  print(f'  - Vocabulary size (V): {cfg.V}')
-  print(f'  - Feed forward dimension (F): {cfg.F}')
+  print(f'  - Model dimension (D): {cfg.model_dim}')
+  print(f'  - Number of heads (H): {cfg.num_heads}')
+  print(f'  - Max sequence length (L): {cfg.seq_len}')
+  print(f'  - Number of layers (N): {cfg.num_layers}')
+  print(f'  - Vocabulary size (V): {cfg.vocab_size}')
+  print(f'  - Feed forward dimension (F): {cfg.expanded_model_dim}')
 
   # Create random input tokens (simulated token IDs)
   rng_key = jax.random.PRNGKey(42)
   input_rng, init_rng = jax.random.split(rng_key)
 
   # Generate random token IDs (integers between 0 and vocab_size-1)
   x_BxL = jax.random.randint(
-    input_rng, shape=(B, L), minval=0, maxval=cfg.V, dtype=jnp.int32
+    input_rng, shape=(B, L), minval=0, maxval=cfg.vocab_size, dtype=jnp.int32
   )
 
   # Initialize model parameters

algoperf/workloads/lm/lm_jax/workload.py

@@ -8,7 +8,7 @@
 from algoperf import jax_sharding_utils, param_utils, spec
 from algoperf.workloads.lm.input_pipeline import get_data_iter
 from algoperf.workloads.lm.lm_jax.nanodo_model import (
-  DoConfig,
+  ModelConfig,
   TransformerDo,
 )
 from algoperf.workloads.lm.workload import BaseLmWorkload
@@ -46,13 +46,13 @@ def init_model_fn(
     aux_dropout_rate: Optional[float] = None,
   ) -> spec.ModelInitState:
     # Initialize NanoDO transformer model
-    cfg = DoConfig(
-      D=self._emb_dim,  # embedding dim
-      H=self._n_heads,  # num heads
-      L=self._seq_len,
-      N=self._n_layers,  # num layers
-      V=self._vocab_size,
-      F=self._mlp_dim,  # feedforward dim
+    cfg = ModelConfig(
+      model_dim=self._emb_dim,  # embedding dim
+      num_heads=self._n_heads,  # num heads
+      seq_len=self._seq_len,
+      num_layers=self._n_layers,  # num layers
+      vocab_size=self._vocab_size,
+      expanded_model_dim=self._mlp_dim,  # feedforward dim
       dtype=jnp.float32,
     )
     self._model = TransformerDo(cfg)

algoperf/workloads/lm/lm_pytorch/plainlm_model.py

@@ -15,15 +15,16 @@
 
 @dataclass
 class ModelConfig:
-  vocab_size: int
+  model_dim: int
+  num_heads: int
   seq_len: int
-  dim: int
-  expand: float
-  n_layers: int
-  n_heads: int
-  rmsnorm_eps: float = 1e-6
-  tie_embeddings: bool = True
+  num_layers: int
+  vocab_size: int
+  expanded_model_dim: int
+  multiple_of: int = 256
+  rmsnorm_epsilon: float = 1e-6
   use_residual_scaling: bool = True
+  tie_embeddings: bool = True
 
 
 class MLP(nn.Module):
@@ -81,13 +82,13 @@ def apply_rotary_emb_complex_like(
 class Attention(nn.Module):
   def __init__(self, cfg: ModelConfig):
     super().__init__()
-    assert cfg.dim % cfg.n_heads == 0
-    self.dim = cfg.dim
-    self.n_heads = cfg.n_heads
-    self.head_dim = cfg.dim // cfg.n_heads
+    assert cfg.model_dim % cfg.num_heads == 0
+    self.dim = cfg.model_dim
+    self.n_heads = cfg.num_heads
+    self.head_dim = cfg.model_dim // cfg.num_heads
 
-    self.w_qkv = nn.Linear(cfg.dim, 3 * cfg.dim, bias=False)
-    self.w_out = nn.Linear(cfg.dim, cfg.dim, bias=False)
+    self.w_qkv = nn.Linear(cfg.model_dim, 3 * cfg.model_dim, bias=False)
+    self.w_out = nn.Linear(cfg.model_dim, cfg.model_dim, bias=False)
     # Split into Q, K, V sections
     wq, wk, wv = torch.chunk(self.w_qkv.weight, 3, dim=0)
     for w in [wq, wk, wv]:
@@ -131,9 +132,9 @@ class Block(nn.Module):
   def __init__(self, layer_id: int, cfg: ModelConfig):
     super().__init__()
     self.attn = Attention(cfg)
-    self.attn_norm = nn.RMSNorm(cfg.dim, eps=cfg.rmsnorm_eps)
-    self.mlp = MLP(dim=cfg.dim, hidden_dim=int(cfg.expand * cfg.dim))
-    self.mlp_norm = nn.RMSNorm(cfg.dim, eps=cfg.rmsnorm_eps)
+    self.attn_norm = nn.RMSNorm(cfg.model_dim, eps=cfg.rmsnorm_epsilon)
+    self.mlp = MLP(dim=cfg.model_dim, hidden_dim=cfg.expanded_model_dim, multiple_of=cfg.multiple_of)
+    self.mlp_norm = nn.RMSNorm(cfg.model_dim, eps=cfg.rmsnorm_epsilon)
     self.layer_id = layer_id
 
   def forward(self, x, freqs_cis):
@@ -144,19 +145,19 @@ def forward(self, x, freqs_cis):
 
 
 class Transformer(nn.Module):
-  def __init__(self, cfg):
+  def __init__(self, cfg: ModelConfig):
     super().__init__()
-    self.n_layers = cfg.n_layers
+    self.n_layers = cfg.num_layers
     self.cfg = cfg
-    head_dim = cfg.dim // cfg.n_heads
-    assert cfg.dim % cfg.n_heads == 0
+    head_dim = cfg.model_dim // cfg.num_heads
+    assert cfg.model_dim % cfg.num_heads == 0
 
-    self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.dim)
+    self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.model_dim)
     self.layers = nn.ModuleList(
-      [Block(idx, cfg) for idx in range(cfg.n_layers)]
+      [Block(idx, cfg) for idx in range(cfg.num_layers)]
     )
-    self.out_norm = nn.RMSNorm(cfg.dim, eps=cfg.rmsnorm_eps)
-    self.lm_head = nn.Linear(cfg.dim, cfg.vocab_size, bias=False)
+    self.out_norm = nn.RMSNorm(cfg.model_dim, eps=cfg.rmsnorm_epsilon)
+    self.lm_head = nn.Linear(cfg.model_dim, cfg.vocab_size, bias=False)
 
     # Initialize freqs_cis on CPU first (more memory efficient)
     self.register_buffer(
@@ -184,7 +185,7 @@ def forward(self, x, targets=None):
     # Make sure we have enough precomputed frequencies
     if L > self.freqs_cis.shape[1]:
       # Need to recompute for longer sequence
-      head_dim = self.cfg.dim // self.cfg.n_heads
+      head_dim = self.cfg.model_dim // self.cfg.num_heads
       new_freqs = precompute_freqs_cis(
         head_dim, max(L, self.cfg.seq_len), 500000
       )
@@ -290,11 +291,11 @@ def main():
   config = ModelConfig(
     vocab_size=50257,  # Common vocab size for tokenizers like BPE or SentencePiece
     seq_len=seq_length,  # Maximum sequence length
-    dim=1024,  # Embedding dimension
-    expand=4.0,  # MLP expansion factor
-    n_layers=12,  # Number of transformer layers
-    n_heads=8,  # Number of attention heads
-    rmsnorm_eps=1e-6,  # RMSNorm epsilon
+    model_dim=1024,  # Embedding dimension
+    expanded_model_dim=4.0,  # MLP expansion factor
+    num_layers=12,  # Number of transformer layers
+    num_heads=8,  # Number of attention heads
+    rmsnorm_epsilon=1e-6,  # RMSNorm epsilon
     tie_embeddings=True,  # Tie embedding and output weights
   )
 

algoperf/workloads/lm/lm_pytorch/workload.py

@@ -39,12 +39,11 @@ def init_model_fn(
     cfg = ModelConfig(
       vocab_size=self._vocab_size,
       seq_len=self._seq_len,
-      dim=self._emb_dim,  # Model dimension
-      expand=self._mlp_dim // self._emb_dim,  # MLP expansion factor
-      # FIXME(rka97): fix expansion factor
-      n_layers=self._n_layers,  # Number of transformer layers
-      n_heads=self._n_heads,  # Number of attention heads
-      rmsnorm_eps=1e-6,
+      model_dim=self._emb_dim,  # Model dimension
+      expanded_model_dim=self._mlp_dim,  # MLP expansion factor
+      num_layers=self._n_layers,  # Number of transformer layers
+      num_heads=self._n_heads,  # Number of attention heads
+      rmsnorm_epsilon=1e-6,
       tie_embeddings=True,
     )
     self._model = Transformer(cfg)