Fix TP, enable test, silence noisy logs (#2761)

Author: ebsmothers

recipes/full_finetune_distributed.py

@@ -231,7 +231,11 @@ def __init__(self, cfg: DictConfig) -> None:
         self._activation_offloading_use_streams = cfg.get(
             "activation_offloading_use_streams", True
         )
-        if self._activation_offloading_use_streams and self.parallel_dims.tp_enabled:
+        if (
+            self._enable_activation_offloading
+            and self._activation_offloading_use_streams
+            and self.parallel_dims.tp_enabled
+        ):
             warn(
                 message=(
                     "Using activation offloading with streams is not advised in tensor parallel, and may "

tests/recipes/test_full_finetune_distributed.py

@@ -11,6 +11,7 @@
 
 import pytest
 import torch
+from packaging import version
 from tests.common import TUNE_PATH
 
 from tests.recipes.utils import (
@@ -130,7 +131,8 @@ def test_loss(
         )
 
     @pytest.mark.skipif(
-        torch.__version__ < "2.8.0", reason="2D parallel test requires PyTorch >= 2.8"
+        version.parse(torch.__version__).base_version < "2.8.0",
+        reason="2D parallel test requires PyTorch >= 2.8",
     )
     @pytest.mark.integration_test
     @pytest.mark.parametrize(

torchtune/models/llama3/_parallelism.py

@@ -37,8 +37,8 @@ def _get_base_llama_tp_training_plan(
         "norm": SequenceParallel(),
         "output": ColwiseParallel(input_layouts=Shard(1), output_layouts=Replicate()),
         "layers.*.attn": layerwise_prepare_module_input_cls(
-            input_layouts=(Shard(1), None),
-            desired_input_layouts=(Replicate(), None),
+            input_layouts=(Shard(1), Shard(1)),
+            desired_input_layouts=(Replicate(), Replicate()),
         ),
         "layers.*.mlp": layerwise_prepare_module_input_cls(
             input_layouts=(Shard(1),),

torchtune/models/llama4/_parallelism.py

@@ -63,8 +63,8 @@ def decoder_only_tp_training_plan(model: nn.Module) -> dict[str, ParallelStyle]:
         layer_plan = {
             f"decoder.layers.{layer_id}.sa_norm": SequenceParallel(),
             f"decoder.layers.{layer_id}.attn": PrepareModuleInput(
-                input_layouts=(Shard(1), None),
-                desired_input_layouts=(Replicate(), None),
+                input_layouts=(Shard(1), Shard(1)),
+                desired_input_layouts=(Replicate(), Replicate()),
             ),
             f"decoder.layers.{layer_id}.attn.q_proj": ColwiseParallel(),
             f"decoder.layers.{layer_id}.attn.k_proj": ColwiseParallel(),