use helpers for selective per layer ac

torchtitan/float8.py

@@ -97,7 +97,6 @@ def convert_to_float8_training(self, model: nn.Module):
             logger.info("Using float8nocompile prototype")
             from torchao.prototype.float8nocompile.float8nocompile_linear_utils import (
                 convert_to_float8_nocompile_training,
-                no_precompute_for_backward_every_nth_layer,
             )
 
             # for full AC or no AC
@@ -175,3 +174,20 @@ def sync_float8_amax_and_scale_history(
         models = [model] if isinstance(model, nn.Module) else model
         for m in models:
             self._sync_float8_amax_and_scale_history(m)
+
+
+def no_precompute_for_backward_every_nth_layer(model: nn.Module, n: int):
+    """Set no_precompute_for_backward to True for every nth layer in the model."""
+    for layer_idx, (layer_id, layer) in enumerate(model.layers.named_children()):
+        if layer_idx % n == 0:
+            logger.info(f"Enabling no_precompute_for_backward for layer {layer_id}")
+            _enable_no_precompute_for_backward(layer)
+
+
+def _enable_no_precompute_for_backward(model: nn.Module):
+    """Recursively set no_precompute_for_backward to True for all linear layers in the given model."""
+    for child_layer in model.children():
+        if isinstance(child_layer, nn.Linear):
+            child_layer.no_precompute_for_backward = True
+        else:
+            _enable_no_precompute_for_backward(child_layer)