Final EKFAC implementation

[LouisYRYJ]

This is it. The final EKFAC merge. Most of it is pretty self-contained and doesn't interfere much with the rest of the code, as we already had set-up much of the work before.

I did slight changes in the code to accomodate for some functionalties that we were missing. In particular, I need to pass a mask tensor to the CovarianceCollector to ensure that we only process the correct activations (and e.g. not the activations of the user tokens).
Technically this needs to happen for any collector, but GradientCollector for example gets around this because gradients will be always 0 at the masked tokens and when we do a*g, the masked token activations get multiplied by 0 and will be ignored (h/t @smarter).

[claude]

Code Review

I've found 5 high-signal issues in this PR:

---

Issue 1: CLAUDE.md violation - torch.cuda.empty_cache() usage

File: bergson/hessians/eigenvectors.py
Lines: 310, 368

The code uses torch.cuda.empty_cache() which violates CLAUDE.md:

torch.cuda.empty_cache() doesn't do what you hope it will do - don't use it.

Fix: Remove both torch.cuda.empty_cache() calls. PyTorch's automatic memory management is sufficient.

---

Issue 2: Bug - Missing gradient masking in kfac.py

File: bergson/hessians/kfac.py
Lines: 68-89

The backward_hook does not apply valid_mask to filter gradients, while forward_hook correctly masks activations. This asymmetry causes:

• Activation covariance computed on [num_valid, I] valid positions
• Gradient covariance computed on [N*S, O] all positions (including padding)

This results in mathematically incorrect KFAC approximation since both covariances must be computed over the same data points.

Fix: Apply the same masking in backward_hook:

def backward_hook(self, module: nn.Module, g: Tensor) -> None:
    """Compute gradient covariance: G^T @ G."""
    name = assert_type(str, module._name)
    S_cov_po = self.S_cov_dict[name]
    mask = self._current_valid_mask
    assert mask is not None, "Valid mask not set for backward hook."
    
    # g: [N, S, O], valid_masks: [N, S] -> select valid positions
    g_bo = g[mask]  # [num_valid, O]
    
    # Compute local covariance
    local_update_oo = g_bo.mT @ g_bo
    # ... rest of the function

---

Issue 3: Bug - Missing gradient masking in tkfac.py

File: bergson/hessians/tkfac.py
Lines: 56-67

Same issue as in kfac.py: backward_hook reshapes gradients to [N*S, O] without masking, while forward_hook masks activations to [num_valid, I]. This creates inconsistent covariance computations.

Fix: Apply the same masking in backward_hook:

def backward_hook(self, module: nn.Module, g: Tensor) -> None:
    """Compute trace-weighted gradient covariance."""
    name = assert_type(str, module._name)
    S_tcov_po = self.S_tcov_dict[name]
    A_tcov_ki = self.A_tcov_dict[name]
    mask = self._current_valid_mask
    assert mask is not None, "Valid mask not set for backward hook."
    
    # g: [N, S, O], valid_masks: [N, S] -> select valid positions
    g_bo = g[mask]  # [num_valid, O]
    a_bi = module._inputs
    # ... rest of the function

---

Issue 4: Bug - Missing loss normalization when use_dataset_labels=True

File: bergson/collector/collector.py
Lines: 626-643

In fwd_bwd_hessian_factory, the use_dataset_labels=True branch is missing per-sample loss normalization:

• When use_dataset_labels=False: losses is reduced to shape [batch_size] via losses.sum(1) / denoms
• When use_dataset_labels=True: losses remains shape [batch_size, seq_len-1] (unreduced)

This inconsistency will cause issues in downstream code expecting per-sample losses.

Fix: Add normalization in the else branch:

else:
    # Use dataset labels instead of sampling
    losses = F.cross_entropy(
        logits[:, :-1].flatten(0, 1),
        y[:, 1:].flatten(0, 1),
        reduction="none",
        ignore_index=-100,
    ).reshape_as(y[:, 1:])
    losses = losses.sum(1) / denoms  # <-- Add this line

---

Issue 5: Bug - dtype string not converted to torch.dtype

File: bergson/hessians/hessian_approximations.py
Lines: 190-208

When hessian_cfg.hessian_dtype != "auto", a string like "bf16" is passed directly as the dtype parameter instead of a torch.dtype object. This will cause a TypeError when torch.zeros(..., dtype=dtype) is called.

Fix: Use convert_precision_to_torch() to convert the string:

hessian_dtype = (
    model.dtype
    if hessian_cfg.hessian_dtype == "auto"
    else convert_precision_to_torch(hessian_cfg.hessian_dtype)
)

Similar pattern is correctly implemented in bergson/score/score.py:284-288.

---

Summary

• 1 CLAUDE.md violation: Remove torch.cuda.empty_cache() calls
• 4 bugs:
  • Missing gradient masking in both KFAC collectors (critical for correctness)
  • Missing loss normalization (causes shape inconsistency)
  • Missing dtype conversion (will cause runtime TypeError)

[LouisYRYJ]

@claude Give it another go. I think i fixed the issues now

[claude]

Claude finished @LouisYRYJ's task in 2m 10s —— View job

---

Re-reviewing PR after fixes

• Read updated files to verify fixes
• Check if CLAUDE.md violation was fixed
• Check if loss normalization was fixed
• Check if dtype conversion was fixed
• Re-evaluate gradient masking issues
• Provide updated review feedback

---

Review Results

Great progress! Most issues have been fixed. Here's the status:

✅ Fixed Issues

1. CLAUDE.md violation - torch.cuda.empty_cache() calls removed from eigenvectors.py:310 and eigenvectors.py:368
2. Loss normalization - Added losses.sum(1) / denoms at collector.py:630
3. Dtype conversion - Now using convert_precision_to_torch() at hessian_approximations.py:196

⚠️ Remaining Concern: Gradient Masking Asymmetry

kfac.py - You mentioned that masked gradients are zero, but I want to clarify a potential issue:

In fwd_bwd_hessian (collector.py:625-630), when use_dataset_labels=False:

losses = F.cross_entropy(
    logits.reshape(-1, logits.size(-1)),
    y[:, 1:].flatten(),  # Contains -100 at masked positions  
    reduction="none",
    # NOTE: No ignore_index=-100 parameter
).reshape_as(y[:, 1:])

Without ignore_index=-100, PyTorch's cross_entropy will:

• Treat -100 as a regular class index
• Compute non-zero loss at those positions
• Generate non-zero gradients flowing backward

This means gradients at masked positions are not automatically zero, creating a mismatch:

• Activation covariance: computed on [num_valid, I] valid positions only (kfac.py:51)
• Gradient covariance: computed on [N*S, O] all positions (kfac.py:74)

The same issue exists in tkfac.py:63 (backward_hook doesn't mask gradients).

Possible resolutions:

1. If the loss computation is correct and gradients ARE zero at masked positions (perhaps through a mechanism I'm not seeing), then tkfac.py should also skip masking for consistency
2. If gradients are NOT zero at masked positions, both kfac.py and tkfac.py should mask gradients in backward_hook to match the activation masking

Could you clarify which case applies? If gradients aren't actually zero at masked positions, I can provide a fix.

---

[LouisYRYJ]

@claude https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.cross_entropy.html ignore_index is set by default to -100, so should be fine?

[claude]

Claude finished @LouisYRYJ's task in 37s —— View job

---

Re-reviewing PR after fixes

• Verify PyTorch's default ignore_index behavior
• Re-examine gradient masking concern
• Update review with corrected information

---

Corrected Review

You're absolutely right! I apologize for the confusion. PyTorch's cross_entropy does indeed have ignore_index=-100 as the default parameter, which means:

• Positions where y == -100 automatically get zero loss
• Gradients flowing backward from those positions are automatically zero
• No explicit masking is needed in the backward hooks

This means:

• kfac.py: The gradient masking concern was invalid. Since masked positions have zero gradients, computing g^T @ g over all positions is mathematically equivalent to computing it over only valid positions.
• tkfac.py: Same reasoning applies - no gradient masking needed.

Final Review Status

✅ All Issues Fixed

1. CLAUDE.md violation - torch.cuda.empty_cache() calls removed from eigenvectors.py
2. Loss normalization - Added losses.sum(1) / denoms at collector.py:644
3. Dtype conversion - Now using convert_precision_to_torch() at hessian_approximations.py:196
4. Gradient masking - Not an issue due to PyTorch's default behavior

The PR looks good to merge! All previously identified issues have been resolved.

---

[luciaquirke]

Follow-up suggestion: could consider having ability to compute the preconditioners + normalizers here too. Seems like a fairly different operation to a full run (less data, more need to parallelize preconditioner computations). Not sure if this makes sense yet

[LouisYRYJ]

I have thought about it too and I agree! I think for now I want to make it lower priority. But sounds like just shuffling some pipelines around. Also we have the compute_eigendecomposition function now in eigenvectors.py which does distirbuted eigenvector decomposition that may be reused

[luciaquirke]

lmk when this is ready!