Add optional gate activation histogram logging during eval

src/levanter/callbacks.py

@@ -47,6 +47,8 @@ def eval_loss_loop(loss_fn, model, dataset, max_batches: Optional[int] = None, n
  load_time = time.time() - time_in
  total_load_time += load_time
  loss = loss_fn(model, batch)
+ if isinstance(loss, tuple):
+ loss, _ = loss
  total_loss += loss.item()
  n += 1
  loss_time = time.time() - time_in - load_time

src/levanter/eval.py

@@ -15,6 +15,7 @@
 from levanter.data import Dataset, ReplicatedBatchLoader
 from levanter.logging import LoadingTimeTrackerIterator
 from levanter.models.lm_model import LmExample, LmHeadModel
+from levanter.tracker.histograms import NBINS
 from levanter.trainer import StepInfo
 from levanter.utils.stat_utils import RunningMean
 from levanter.utils.tree_utils import inference_mode
@@ -34,6 +35,7 @@ class EvalResult:
  tag_macro_losses: dict[str, float] # per tag average-per-token loss
  tag_micro_losses: dict[str, float] # per tag total loss, for "parent" tags
  total_eval_loading_time: float
+ extras: dict[str, float]
 
 
 class DomainTaggedDataset(Dataset[tuple[T, hax.NamedArray]]):
@@ -123,6 +125,17 @@ def eval_callback(step: StepInfo):
  _join_prefix(prefix, "loading_time"): result.total_eval_loading_time,
  _join_prefix(prefix, "total_time"): time_fn(),
  }
+ if (gate_hist := result.extras.get("gate_hist", None)) is not None:
+ pos_idx = NBINS // 2 + 1
+ log_dict[_join_prefix(prefix, "gate_hist/all")] = np.array(gate_hist.sum(axis=0))
+ num_gt0 = gate_hist[:, pos_idx:].sum().item()
+ total = gate_hist.sum().item()
+ log_dict[_join_prefix(prefix, "gate_gt0/all")] = num_gt0 / total
+ for i in range(gate_hist.shape[0]):
+ log_dict[_join_prefix(prefix, f"gate_hist/layer{i+1}")] = np.array(gate_hist[i])
+ num_gt0 = gate_hist[i, pos_idx:].sum().item()
+ total = gate_hist[i].sum().item()
+ log_dict[_join_prefix(prefix, f"gate_gt0/layer{i+1}")] = num_gt0 / total
 
  logger.info(f"{prefix} loss: {result.micro_avg_loss:.3f}")
  for tag, loss in result.tag_macro_losses.items():
@@ -185,12 +198,12 @@ def __init__(
 
  @hax.named_jit(out_axis_resources=axis_mapping)
  def accum_for_batch(
- m: LmHeadModel, state: tuple[RunningMean, RunningMean], batch: LmExample, tags: hax.NamedArray
+ m: LmHeadModel, state: tuple[RunningMean, RunningMean, dict], batch: LmExample, tags: hax.NamedArray
  ):
  m = inference_mode(m, True)
  with hax.axis_mapping(axis_mapping):
- total_mean, mean_per_tag = state
- losses = m.compute_loss(batch, reduction=None, reduction_axis=())
+ total_mean, mean_per_tag, total_extras = state
+ losses, extras = m.compute_loss(batch, reduction=None, reduction_axis=())
  mask = batch.loss_mask # [Batch, Token]
  this_tokens = hax.einsum("->", mask)
  this_loss = hax.einsum("->", losses, mask) # to scalar
@@ -203,23 +216,32 @@ def accum_for_batch(
  safe_mean = hax.where(this_tokens_per_tag, this_loss_per_tag / this_tokens_per_tag, 0.0)
  mean_per_tag = mean_per_tag.add(safe_mean, this_tokens_per_tag)
 
- return mean, mean_per_tag
+ if extras:
+ for key in extras:
+ curr = total_extras.get(key, jnp.zeros_like(extras[key]))
+ total_extras[key] = extras[key] + curr
+
+ return mean, mean_per_tag, total_extras
 
  self.accum_for_batch = accum_for_batch
 
  def evaluate(self, m: LmHeadModel):
  total_loss = jnp.zeros(())
  mean_losses_per_tag = hax.zeros(self.dataset.Tag, dtype=np.float32)
 
- state = (RunningMean.zeros_like(total_loss), RunningMean.zeros_like(mean_losses_per_tag))
+ state: tuple[RunningMean, RunningMean, dict] = (
+ RunningMean.zeros_like(total_loss),
+ RunningMean.zeros_like(mean_losses_per_tag),
+ {},
+ )
  state = hax.shard(state)
 
  iterator = LoadingTimeTrackerIterator(self.loader)
 
  for batch, tags in tqdm.tqdm(iterator, "eval"):
  state = self.accum_for_batch(m, state, batch, tags)
 
- total_loss, losses_per_tag = state
+ total_loss, losses_per_tag, extras = state
 
  micro_avg_loss = total_loss.mean.item()
  tag_avg_loss = losses_per_tag.mean
@@ -252,4 +274,4 @@ def evaluate(self, m: LmHeadModel):
  tag_micro_loss[tag] = mean_loss_per_tag_cpu[index]
  # no macro loss for the leaf tags
 
- return EvalResult(micro_avg_loss, macro_avg_loss, tag_macro_loss, tag_micro_loss, iterator.total_time)
+ return EvalResult(micro_avg_loss, macro_avg_loss, tag_macro_loss, tag_micro_loss, iterator.total_time, extras)

src/levanter/main/viz_logprobs.py

@@ -72,7 +72,7 @@ def main(config: VizGpt2Config):
  def compute_log_probs(model: LmHeadModel, example: LmExample):
  model = inference_mode(model, True)
  model = mp.cast_to_compute(model)
- logprobs = model.compute_loss(example, reduction=None)
+ logprobs, _ = model.compute_loss(example, reduction=None)
  # roll forward to get the loss for each predicted token
  logprobs = hax.roll(logprobs, 1, Pos)
  return logprobs.rearrange((EvalBatch, Pos)).array

src/levanter/models/asr_model.py

@@ -105,15 +105,15 @@ def compute_loss(
  across the reduction axis (with reduction_axis=None meaning all axes). If reduction is None, the loss is not
  reduced, and the result is a named array with axes (*batch axes, sequence_length).
  """
- logits = self(example.audio, example.tokens, example.attn_mask, key=key)
+ logits, extras = self(example.audio, example.tokens, example.attn_mask, key=key)
  logits = logits.astype(jnp.float32)
  targets = hax.roll(example.tokens, -1, axis=self.Pos.name)
  target_y = hax.nn.one_hot(targets, self.Vocab, dtype=logits.dtype)
  loss = cross_entropy_loss(
  logits, self.Vocab, target_y, reduction, reduction_axis=reduction_axis, where=example.loss_mask
  )
 
- return loss
+ return loss, extras
 
  @property
  def vocab_size(self) -> int:

src/levanter/models/gemma.py

@@ -30,6 +30,7 @@
 from levanter.models.lm_model import LmConfig, LmHeadModel
 from levanter.types import BlockFoldable
 from levanter.utils.flop_utils import lm_flops_per_token
+from levanter.utils.py_utils import cached_classproperty
 
 
 silence_transformer_nag()
@@ -75,7 +76,6 @@ class GemmaConfig(HFCompatConfig):
  vocab_size: int = 256_000
  num_layers: int = 18
  num_heads: int = 8
- head_dim: int = 256
  num_kv_heads: int = 1
  attn_dropout = 0.0
  norm_eps = 1e-6
@@ -106,10 +106,14 @@ class GemmaConfig(HFCompatConfig):
  Mlp = property(lambda self: Axis(name="mlp", size=self.intermediate_dim))
  HeadSize = property(lambda self: Axis(name="head_size", size=self.hidden_dim // self.num_heads))
 
+ @property
+ def head_dim(self) -> int: return self.hidden_dim // self.num_heads
+
  def __post_init__(self):
  assert (
  self.num_heads % self.num_kv_heads == 0
  ), f"num_heads={self.num_heads} not divisible by num_kv_heads={self.num_kv_heads}."
+ assert (self.head_dim * self.num_heads) == self.hidden_dim, "head_dim * num_heads must equal hidden_dim."
 
  def hf_checkpoint_converter(self) -> HFCheckpointConverter["GemmaConfig"]: # type: ignore
  return HFCheckpointConverter(
@@ -129,7 +133,9 @@ def from_hf_config(cls, hf_config: HfConfig):
  if hf_config.hidden_activation:
  activation_function = hf_config.hidden_activation
  else:
- activation_function = hf_config.hidden_act
+ # This is the implementation in huggingface
+ # https://github.com/huggingface/transformers/blob/12b1620e615592fbf099d4ec44af7b9f2d1b48aa/src/transformers/models/gemma/modeling_gemma.py#L200
+ activation_function = "gelu_pytorch_tanh"
 
  if activation_function == "gelu_pytorch_tanh":
  activation_function = "gelu_new"
@@ -168,7 +174,7 @@ def to_hf_config(self, vocab_size: int, config_overrides: Optional[Dict] = None)
  num_hidden_layers=self.num_layers,
  num_attention_heads=self.num_heads,
  num_key_value_heads=self.num_kv_heads,
- head_dim=self.hidden_dim // self.num_heads,
+ head_dim=self.head_dim,
  hidden_activation=(
  "gelu_pytorch_tanh" if self.activation_function == "gelu_new" else self.activation_function
  ),
@@ -263,9 +269,9 @@ def __call__(self, x: NamedArray, mask: Optional[NamedArray | AttentionMask], *,
  # MLP and skip connection
  residual = x
  x = self.post_attention_layernorm(x)
- mlp_output = self.mlp(x, key=k_mlp)
+ mlp_output, extras = self.mlp(x, key=k_mlp)
  output = residual + mlp_output
- return output
+ return output, extras
 
 
 class GemmaTransformer(StateDictSerializationMixin, eqx.Module):
@@ -292,10 +298,10 @@ def init(config: GemmaConfig, *, key) -> "GemmaTransformer":
  @named_call
  def __call__(self, x: NamedArray, attn_mask: Optional[NamedArray | AttentionMask], *, key) -> NamedArray:
  keys = maybe_rng_split(key, self.config.num_layers) if key is not None else None
- x = self.layers.fold(x, mask=attn_mask, key=keys)
+ x, extras = self.layers.scan(x, mask=attn_mask, key=keys)
  x = self.norm(x)
 
- return x
+ return x, extras
 
  def from_state_dict(self, state_dict: StateDict, prefix: Optional[str] = None):
  if isinstance(self.layers, Stacked):
@@ -358,9 +364,9 @@ def __call__(
  The attn_mask from training pipeline may be an AttentionMask object instead of NamedArray
  """
  x = self.embeddings.embed(input_ids)
- x = self.transformer(x, attn_mask=attn_mask, key=key)
+ x, extras = self.transformer(x, attn_mask=attn_mask, key=key)
  lm_logits = self.embeddings.unembed(x)
- return lm_logits
+ return lm_logits, extras
 
  def resize_vocab(self, new_size: int, key=None) -> "LmHeadModel[GemmaConfig]":
  new_embeddings = self.embeddings.resize_embeddings(new_size, key=key)

src/levanter/models/gpt2.py

@@ -29,6 +29,7 @@
 from levanter.models.attention import AttentionBackend, AttentionMask, dot_product_attention
 from levanter.models.lm_model import LmConfig
 from levanter.utils.flop_utils import lm_flops_per_token
+from levanter.utils.py_utils import cached_classproperty
 
 
 silence_transformer_nag()

src/levanter/models/llama.py

@@ -30,6 +30,8 @@
 from levanter.models.lm_model import LmConfig, LmHeadModel
 from levanter.types import BlockFoldable
 from levanter.utils.flop_utils import lm_flops_per_token
+from levanter.utils.py_utils import cached_classproperty
+from levanter.tracker.histograms import get_bins, sharded_histogram
 
 
 silence_transformer_nag()
@@ -78,6 +80,7 @@ class LlamaConfig(HFCompatConfig):
  use_bias: bool = False
  use_layer_norm_weight: bool = True
  rope_scaling: Optional[dict] = None
+ measure_act_stats: bool = True
 
  reference_checkpoint: str = "meta-llama/Llama-2-7b-hf"
  tokenizer: Optional[str] = None
@@ -181,10 +184,17 @@ class LlamaMlp(eqx.Module, StateDictSerializationMixin):
  up_proj: hnn.Linear # projection from Embed to Mlp
  down_proj: hnn.Linear # projection from Mlp to Embed
  act: Callable = eqx.static_field()
+ measure_act_stats: bool = False
 
  @staticmethod
  def init(
- Embed: Axis, Mlp: Axis, activation_fn: Union[str, Callable], *, key, use_bias: bool = False
+ Embed: Axis,
+ Mlp: Axis,
+ activation_fn: Union[str, Callable],
+ *,
+ key,
+ use_bias: bool = False,
+ measure_act_stats=False,
  ) -> "LlamaMlp":
  k_fc, k_up_proj, k_down_proj = jrandom.split(key, 3)
  gate_proj = hnn.Linear.init(Out=Mlp, In=Embed, key=k_fc, use_bias=use_bias, out_first=True)
@@ -193,16 +203,20 @@ def init(
  if isinstance(activation_fn, str):
  activation_fn = ACT2FN[activation_fn]
  act = activation_fn # type: ignore
- return LlamaMlp(gate_proj, up_proj, down_proj, act)
+ get_bins() # initialize bins
+ return LlamaMlp(gate_proj, up_proj, down_proj, act, measure_act_stats)
 
  @named_call
  def __call__(self, x: NamedArray, *, key=None) -> NamedArray:
  k_gate, k_up, k_down = maybe_rng_split(key, 3)
  hidden_states = self.gate_proj(x, key=k_gate)
+ extras = {}
+ if self.measure_act_stats:
+ extras["gate_hist"] = sharded_histogram(hidden_states.array, bins=get_bins())
  hidden_states = self.act(hidden_states)
  hidden_states = hidden_states * self.up_proj(x, key=k_up)
  outputs = self.down_proj(hidden_states, key=k_down)
- return outputs
+ return outputs, extras
 
  def from_state_dict(self, state_dict: StateDict, prefix: Optional[str] = None):
  # unflatten the linear layers of HF state_dict to match the shape of LlamaMlp
@@ -402,6 +416,7 @@ def init(config: LlamaConfig, *, key) -> "LlamaDecoderLayer":
  config.activation_function,
  key=k_mlp,
  use_bias=config.use_bias,
+ measure_act_stats=config.measure_act_stats,
  )
  ln_1 = config.mk_LayerNorm(config.Embed)
  ln_2 = config.mk_LayerNorm(config.Embed)
@@ -420,9 +435,9 @@ def __call__(self, x: NamedArray, mask: Optional[NamedArray | AttentionMask], *,
  # MLP and skip connection
  residual = x
  x = self.post_attention_layernorm(x)
- mlp_output = self.mlp(x, key=k_mlp)
+ mlp_output, extras = self.mlp(x, key=k_mlp)
  output = residual + mlp_output
- return output
+ return output, extras
 
 
 class LlamaTransformer(StateDictSerializationMixin, eqx.Module):
@@ -449,10 +464,10 @@ def init(config: LlamaConfig, *, key) -> "LlamaTransformer":
  @named_call
  def __call__(self, x: NamedArray, attn_mask: Optional[NamedArray | AttentionMask], *, key) -> NamedArray:
  keys = maybe_rng_split(key, self.config.num_layers) if key is not None else None
- x = self.layers.fold(x, mask=attn_mask, key=keys)
+ x, extras = self.layers.scan(x, mask=attn_mask, key=keys)
  x = self.norm(x)
 
- return x
+ return x, extras
 
  def from_state_dict(self, state_dict: StateDict, prefix: Optional[str] = None):
  if isinstance(self.layers, Stacked):
@@ -544,9 +559,9 @@ def __call__(
  """
  k_t, k_head = maybe_rng_split(key, 2)
  x = self.embeddings.embed(input_ids)
- x = self.transformer(x, attn_mask=attn_mask, key=k_t)
+ x, extras = self.transformer(x, attn_mask=attn_mask, key=k_t)
  lm_logits = self.lm_head(x, key=k_head)
- return lm_logits
+ return lm_logits, extras
 
  def resize_vocab(self, new_size: int, key=None) -> "LmHeadModel[LlamaConfig]":
  new_Vocab = self.Vocab.resize(new_size)

src/levanter/models/lm_model.py

@@ -127,6 +127,10 @@ def compute_loss(
  reduced, and the result is a named array with axes (*batch axes, sequence_length).
  """
  logits = self(example.tokens, example.attn_mask, key=key)
+ extras = None
+ if isinstance(logits, tuple):
+ assert len(logits) == 2
+ logits, extras = logits
  # TODO: would be nice if we made the dtype configurable
  logits = logits.astype(jnp.float32)
  targets = hax.roll(example.tokens, -1, axis=self.Pos.name)
@@ -135,7 +139,7 @@ def compute_loss(
  logits, self.Vocab, target_y, reduction, reduction_axis=reduction_axis, where=example.loss_mask
  )
 
- return loss
+ return loss, extras
 
  @property
  def vocab_size(self) -> int:

src/levanter/models/mistral.py

@@ -23,6 +23,7 @@
 from levanter.models.llama import LlamaConfig, LlamaEmbedding, LlamaTransformer
 from levanter.models.lm_model import LmConfig, LmHeadModel
 from levanter.utils.flop_utils import lm_flops_per_token
+from levanter.utils.py_utils import cached_classproperty
 
 
 silence_transformer_nag()
@@ -192,9 +193,9 @@ def __call__(
  """
  k_t, k_head = maybe_rng_split(key, 2)
  x = self.embeddings.embed(input_ids)
- x = self.transformer(x, attn_mask=attn_mask, key=k_t)
+ x, extras = self.transformer(x, attn_mask=attn_mask, key=k_t)
  lm_logits = self.lm_head(x, key=k_head)
- return lm_logits
+ return lm_logits, extras
 
  def resize_vocab(self, new_size: int, key=None) -> "LmHeadModel[MistralConfig]":
  new_Vocab = self.Vocab.resize(new_size)