TRF034

Layer classes held in an nn.ModuleList must subclass GradientCheckpointingLayer.

   
Default Enabled
Scope Models added on or after 2026-06-20
Source mlinter/trf034.py
Show in terminal mlinter --rule TRF034

What it does

In modeling_*.py and modular_*.py, flags a locally-defined *Layer/*Block class instantiated in an nn.ModuleList(...) that does not reach GradientCheckpointingLayer through its base chain; modular files follow relative imports into sibling models, and unresolved chains are inconclusive. Out of scope: a model that never sets supports_gradient_checkpointing = True, which raises from gradient_checkpointing_enable() instead of skipping a layer; a layer holding nn.BatchNorm*/nn.InstanceNorm*, whose statistics would be recomputed twice; and a stack that is not the model’s token mixer, shown by an attention, modulation, mixer or SSM module assigned as self.x = Y(...).

Why is this bad?

gradient_checkpointing_enable() wraps a layer only if it is a GradientCheckpointingLayer. A plain nn.Module on the trunk is skipped silently, so training looks checkpointed while still allocating full activations, and the OOM surfaces far from the cause. Elsewhere – a conv backbone, a decode head – the trade is the author’s call, so the rule stays out.

Example

-class AcmeDecoderLayer(nn.Module):
+class AcmeDecoderLayer(GradientCheckpointingLayer):
     def __init__(self, config, layer_idx):
         super().__init__()

Suppressing this rule

Add a # trf-ignore: TRF034 comment on the flagged line or the line directly above it. See Suppressing rules for whole-file directives and when a suppression is the wrong answer.

Allowlisted models

1 model are exempt from TRF034 in mlinter/rules.toml, because they predate the convention and cannot be changed without breaking backward compatibility.

Show the 1 allowlisted model
  • x_clip