remove unused FSDP float8 all-gather integration (#182)

Summary:
All of the removed code is outdated, we will try again with per-parameter sharding FSDP.

Pull Request resolved: #182

Test Plan:
```
./test/test_everything.sh
```

Reviewed By: awgu

Differential Revision: D52648709

Pulled By: vkuzo

fbshipit-source-id: 202c2675e4e96a035d5f345346f1a1a3816f02ae

Author: vkuzo

float8_experimental/dynamic_linear/dynamic_float8_linear.py

@@ -46,16 +46,9 @@ class Float8DynamicLinear(torch.nn.Linear):
     conversion to fp8 of the input and weight tensors.
     """
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.add_weight_tag()
-
     def forward(self, x):
         x_fp8 = self.cast_to_float8(x)
-        if getattr(self, "_w_fp8", None) is not None:  # FSDP handled the cast
-            w_fp8 = self._w_fp8
-        else:
-            w_fp8 = self.cast_to_float8(self.weight)
+        w_fp8 = self.cast_to_float8(self.weight)
 
         y = torch.nn.functional.linear(x_fp8, w_fp8, self.bias)
 
@@ -64,11 +57,6 @@ def forward(self, x):
 
         return y
 
-    def add_weight_tag(self):
-        # We add a tag to the weight nn.Parameter in order to signal
-        # To FSDP that this param is a weight
-        self.weight._is_fp8_weight = True
-
     def cast_to_float8(self, inpt_tensor):
         scale = tensor_to_scale(inpt_tensor, torch.float8_e4m3fn)
         return Float8Tensor.to_float8(
@@ -92,5 +80,4 @@ def from_float(cls, mod, emulate: bool = False):
         new_mod.weight = mod.weight
         new_mod.bias = mod.bias
         new_mod.emulate = emulate
-        new_mod.add_weight_tag()
         return new_mod

float8_experimental/float8_linear.py

@@ -171,13 +171,6 @@ def __init__(self, *args, **kwargs):
         # Note: this is not used in non-TP code.
         self.use_sequence_parallel = False
 
-        # Save the Float8Tensor constructor for FSDP.
-        # N.B. Do not partially apply the scale into the constructor because
-        # buffer Python IDs are not preserved by `nn.Module.to()` and the
-        # module could be moved to GPU after this constructor. Instead, FSDP
-        # will access the scale when it has ensured that it is on GPU.
-        self._float8_tensor_ctor = lambda *args, **kwargs: Float8Tensor(*args, **kwargs)
-
         # pre_forward and post_forward are currently broken with FSDP
         # and torch.compile, this option can disable them
         self.enable_pre_and_post_forward = config.enable_pre_and_post_forward
@@ -312,30 +305,18 @@ def float8_post_forward(self):
         self.is_amax_initialized = True
         self.amax_and_scale_synced = False
 
-    def add_weight_tag(self):
-        # We add a tag to the weight nn.Parameter in order to signal
-        # To FSDP that this param is a weight
-        self.weight._is_fp8_weight = True
-
 
 class Float8Linear(Float8LinearMixin, torch.nn.Linear):
     """
     A wrapper around a `torch.nn.Linear` module which does fp8 compute, and tracks
     scales in way friendly to delayed scaling.
     """
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.add_weight_tag()
-
     def forward(self, x):
         self.float8_pre_forward(x)
 
         x_fp8 = self.cast_x_to_float8(x, self.is_amax_initialized)
-        if getattr(self, "_w_fp8", None) is not None:  # FSDP handled the cast
-            w_fp8 = self._w_fp8
-        else:
-            w_fp8 = self.cast_w_to_float8(self.weight, self.is_amax_initialized)
+        w_fp8 = self.cast_w_to_float8(self.weight, self.is_amax_initialized)
 
         y = torch.matmul(x_fp8, w_fp8.t())
 
@@ -366,5 +347,4 @@ def from_float(cls, mod, emulate: bool = False):
         new_mod.emulate = emulate
         # I think its okay to send all params and buffers to device
         new_mod.to(mod.weight.device)
-        new_mod.add_weight_tag()
         return new_mod

float8_experimental/tp_linear.py

@@ -49,10 +49,7 @@ def forward(self, input_: torch.Tensor) -> torch.Tensor:  # type: ignore
                 input_parallel, self.is_amax_initialized
             )
 
-        if getattr(self, "_w_fp8", None) is not None:  # FSDP handled the cast
-            w_fp8 = self._w_fp8
-        else:
-            w_fp8 = self.cast_w_to_float8(self.weight, self.is_amax_initialized)
+        w_fp8 = self.cast_w_to_float8(self.weight, self.is_amax_initialized)
 
         # Matrix multiply.
         output_parallel = torch.matmul(input_parallel_fp8, w_fp8.t())
@@ -101,7 +98,6 @@ def from_float(cls, mod, emulate=False):
         device_to_use = next(mod.parameters()).device
         new_mod.to(device_to_use)
         new_mod.emulate = emulate
-        new_mod.add_weight_tag()
         # TODO: test when creation is on cuda
         return new_mod
 
@@ -137,10 +133,7 @@ def forward(self, input_: torch.Tensor) -> torch.Tensor:  # type:ignore
             input_parallel, self.is_amax_initialized
         )
 
-        if getattr(self, "_w_fp8", None) is not None:  # FSDP handled the cast
-            w_fp8 = self._w_fp8
-        else:
-            w_fp8 = self.cast_w_to_float8(self.weight, self.is_amax_initialized)
+        w_fp8 = self.cast_w_to_float8(self.weight, self.is_amax_initialized)
 
         # Matrix multiply.
         output_parallel = torch.matmul(input_parallel_fp8, w_fp8.t())
@@ -196,7 +189,6 @@ def from_float(cls, mod, emulate=False):
         device_to_use = next(mod.parameters()).device
         new_mod.to(device_to_use)
         new_mod.emulate = emulate
-        new_mod.add_weight_tag()
         return new_mod
 
 

test/test_base.py

@@ -175,11 +175,6 @@ def test_linear_bias(
             y.dtype == torch.bfloat16
         ), f"y.dtype is {y.dtype}, expected {torch.bfloat16}"
 
-    def test_linear_float8_weight_tag(self) -> None:
-        m_ref = nn.Linear(16, 32, bias=False, device="cuda")
-        m_fp8 = Float8Linear.from_float(copy.deepcopy(m_ref))
-        assert m_fp8.weight._is_fp8_weight
-
     @pytest.mark.parametrize("linear_type", [LinearType.DELAYED, LinearType.DYNAMIC])
     @pytest.mark.parametrize(
         "linear_dtype", [torch.float16, torch.bfloat16, torch.float32]