Adding back comments

Author: parambole

MaxText/layers/models.py

@@ -180,12 +180,12 @@ def __call__(self, hidden_states: jnp.ndarray, deterministic: bool, model_mode:
           inputs_shape=y.shape,
           features=cfg.vocab_size,
           weight_dtype=cfg.weight_dtype,
-          dtype=jnp.float32 if cfg.logits_dot_in_fp32 else cfg.dtype,
+          dtype=jnp.float32 if cfg.logits_dot_in_fp32 else cfg.dtype,  # for logit training stability
           kernel_axes=("embed", "vocab"),
           name="logits_dense",
           matmul_precision=self.config.matmul_precision,
       )
-      # Then, call the instance with the input tensor.
+      # We do not quantize the logits matmul.
       logits = dense_layer(y)
 
     # 4. Final Casting