Defer non-differentiable checks to attribution. (#1826)

[styusuf]

Summary:

Previously, LayerGradientAttributor._find_differentiable_layers ran an upfront torch.autograd.grad call against the first task to identify which layers are differentiable, filtering them out before attribution. This had two problems:

1. A layer differentiable for the first task but non-differentiable for other tasks would pass the check, then crash during attribution with a None gradient error in _reduce_list. There was a patch fix for this here: D102939482
2. A layer non-differentiable for the first task but differentiable for other tasks would be discarded entirely, losing valid attributions for those tasks.

This diff removes the upfront differentiability check and instead defers detection to the attribution computation itself. The key insight is that differentiability is a per-task property, not a per-layer property — a layer may be connected to the computation graph for some tasks and not others.

Non-differentiable vs. zero attribution

These are semantically different:

• Non-differentiable: The layer has no path in the computation graph to the task's output. torch.autograd.grad returns None. The gradient literally does not exist — the layer plays no role in producing that task's prediction.
• Zero attribution: The layer IS connected to the task's output and autograd.grad returns a real tensor, but the gradient values happen to be zero (or near-zero). The layer participates in the computation but its contribution cancels out or is negligible for the given input.

Both produce zero-valued tensors in the final output. The distinction is preserved through None propagation in the gradient pipeline and surfaced via logging in _consolidate_task_attributions, which logs layers that are non-differentiable for ALL tasks. For layers non-differentiable for only some tasks, the None entries are replaced with zero tensors to allow stacking — these are indistinguishable from genuine zero attributions in the output tensor.

Changes

gradient.py — compute_layer_gradients_and_eval:

• Replaces D102939482's bulk None→zeros substitution with per-layer handling
• When allow_unused=True and ALL grads for a layer are None, returns None for that layer (preserves non-differentiable signal for downstream consumers)
• When only SOME grads are None (partial differentiability edge case), substitutes zeros to protect _reduce_list's Tensor-only invariant
• None return is gated behind allow_unused=True so existing callers that don't pass allow_unused are unaffected

layer_gradient_x_activation.py — LayerGradientXActivation.attribute:

• Multi-layer path: propagates None grads through to the caller instead of crashing
• multiply_gradient_acts: guards against individual None gradient elements

layer_attr.py — LayerGradientAttributor:

• Renames _find_differentiable_layers → _approve_supported_layers: approves all layers with valid output format (tensor, tuple of tensors, floating-point, requires_grad) without checking differentiability
• _get_attributions: passes allow_unused=True to autograd.grad, allowing None returns for non-differentiable layers
• _consolidate_task_attributions: handles per-task None attributions — all-None layers are logged as non-differentiable and return None; partial-None layers are zero-filled from a reference attribution shape
• Fixes _log_layers_not_differentiable type signature to match 3-tuple call site

common.py — _reduce_list: handles None values that may appear when allow_unused=True is used

Differential Revision: D103563408

[meta-codesync]

@styusuf has exported this pull request. If you are a Meta employee, you can view the originating Diff in D103563408.