From 73ce4b419cb6eba6b040eaeef51834e7224636d9 Mon Sep 17 00:00:00 2001
From: ngoyal2707 <ngoyal2707@users.noreply.github.com>
Date: Wed, 24 Jan 2024 15:49:34 -0500
Subject: [PATCH 01/20] added option for no PG validation for faster init
 (#1161)

Co-authored-by: Naman Goyal <naman@fb.com>
---
 fairscale/nn/data_parallel/fully_sharded_data_parallel.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index 7d7fce9a4..e5cf25660 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -368,6 +368,7 @@ def __init__(
         gradient_predivide_factor: Optional[float] = None,
         limit_all_gather_events: bool = False,
         limit_reduce_scatter_events: bool = False,
+        should_validate_process_group: bool = True,
     ):
         try:
             import torch._C
@@ -451,7 +452,7 @@ def __init__(
             raise ValueError(f"offload type: '{offload_config.offload_type}' requires flatten_parameters=True")
 
         # skip validation if the process group was created above
-        if process_group:
+        if process_group and should_validate_process_group:
             validate_process_group(self.compute_device, self.process_group)
 
         # enable pytorch sync_bn just in case model contains sync_bn layers.

From 33457b38363bcc7a19975a59c97bb72188265b34 Mon Sep 17 00:00:00 2001
From: Shikai Li <shikaili@fb.com>
Date: Tue, 12 Mar 2024 19:52:59 -0700
Subject: [PATCH 02/20] Mirros Jiecao's change.

---
 .../fully_sharded_data_parallel.py            | 229 ++++++++++++++++--
 fairscale/nn/misc/__init__.py                 |   2 +-
 fairscale/nn/misc/flatten_params_wrapper.py   |   3 +-
 3 files changed, 206 insertions(+), 28 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index e5cf25660..b66c09bd8 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -42,8 +42,11 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from torch.nn.parameter import Parameter
+import transformer_engine.pytorch as te
+from transformer_engine.pytorch.cpp_extensions import cast_to_fp8, DType, FP8FwdTensors
+from transformer_engine.pytorch.fp8 import amax_and_scale_update, FP8GlobalStateManager
 
-from fairscale.nn.misc import FlattenParamsWrapper
+from fairscale.nn.misc import FlatParameter, FlattenParamsWrapper
 from fairscale.nn.wrap import auto_wrap, config_auto_wrap_policy, enable_wrap
 from fairscale.utils.containers import apply_to_tensors
 from fairscale.utils.parallel import (
@@ -150,6 +153,14 @@ class OffloadConfig:
     dir: Optional[str] = None
 
 
+def _is_fp8_dtype(dtype: torch.dtype) -> bool:
+    return dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
+
+
+def _is_te_module_with_weights(m: nn.Module) -> bool:
+    return isinstance(m, (te.Linear, te.LayerNormLinear, te.LayerNormMLP))
+
+
 class FullyShardedDataParallel(nn.Module):
     """
     A wrapper for sharding Module parameters across data parallel workers. This
@@ -369,6 +380,7 @@ def __init__(
         limit_all_gather_events: bool = False,
         limit_reduce_scatter_events: bool = False,
         should_validate_process_group: bool = True,
+        fp8_allgather: bool = False,
     ):
         try:
             import torch._C
@@ -434,6 +446,7 @@ def __init__(
         self.force_input_to_fp32 = force_input_to_fp32
         self.verbose = verbose
         self.state_dict_on_rank_0_only = state_dict_on_rank_0_only
+        self.fp8_allgather = fp8_allgather
         # Experimental feature for now. Use at your own risk.
         self.ssd_offload = True if offload_config and offload_config.offload_type == "ssd_offload" else False
 
@@ -488,9 +501,22 @@ def __init__(
         non_flatten_params = params
         param_name_groups = [[n] for n in param_names]
         if self.flatten_parameters:
-            to_be_flatten_params = [params]
-            non_flatten_params = []
-            param_name_groups = [param_names]
+            to_be_flatten_params = [
+                [
+                    params[i]
+                    for i in range(len(params))
+                    if "norm_weight" not in param_names[i]
+                ]
+            ]
+            non_flatten_params = [
+                params[i]
+                for i in range(len(params))
+                if "norm_weight" in param_names[i]
+            ]
+            param_name_groups = [
+                [n for n in param_names if "norm_weight" not in n],
+                [n for n in param_names if "norm_weight" in n],
+            ]
         del param_names
 
         self._fsdp_wrapped_module: nn.Module = FlattenParamsWrapper(
@@ -559,6 +585,10 @@ def __init__(
         self._all_gather_free_event_queue = _FreeEventQueue() if limit_all_gather_events else None
         self._reduce_scatter_free_event_queue = _FreeEventQueue() if limit_reduce_scatter_events else None
 
+    @property
+    def _is_fp8_compute(self) -> bool:
+        return _is_fp8_dtype(self.compute_dtype)
+
     def _get_gradient_predivide_factor(self, world_size: int) -> float:
         factor: int = 1
         while world_size % factor == 0 and world_size / factor > factor:
@@ -688,7 +718,11 @@ def _cast_buffers(
     @property
     def params_with_grad(self) -> List[Parameter]:
         """[p for p in self.parameters() if p.grad is not None]"""
-        return [p for p in self.parameters() if (p.grad is not None or p.main_grad is not None)]
+        return [
+            p
+            for p in self.parameters()
+            if (p.grad is not None or getattr(p, "main_grad", None) is not None)
+        ]
 
     @torch.no_grad()
     def clip_grad_norm_(
@@ -791,7 +825,9 @@ def _shard_parameters_(self) -> None:
                 assert p.dtype == torch.float32
 
             # If world_size is 1, then we all-reduce grads instead of sharding.
-            p._is_sharded = self.world_size > 1
+            p._is_sharded = (self.world_size > 1) and (
+                not self._is_fp8_compute or isinstance(p, FlatParameter)
+            )
             p._orig_size = p.data.size()
 
             if not p._is_sharded:
@@ -1175,7 +1211,7 @@ def summon_full_params(self, recurse: bool = True, volatile: bool = False) -> Ge
                         non_shared_params
                     ), f"{len(full_tensors)} vs. {len(non_shared_params)}"
                     for p, (full_tensor, safe_to_free) in zip(non_shared_params, full_tensors):
-                        if not volatile:
+                        if not volatile and p._is_sharded:
                             # Copy any changes made to the full params back into
                             # the corresponding local shards.
                             local_shard, _ = self._get_shard(full_tensor)
@@ -1237,6 +1273,20 @@ def _lazy_init(self) -> None:
             # ``optim.step()`` is done before we all-gather parameters.
             self._wait_for_previous_optim_step()
 
+    def _shard_dtype(self, p: Parameter) -> torch.dtype:
+        """
+        Return the dtype to use for the sharded parameters.
+
+        Returns:
+            The dtype to use for the sharded parameters.
+        """
+        if self._is_fp8_compute and not isinstance(p, FlatParameter):
+            # Assume non flattened are precision critical like norm
+            assert not p._is_sharded
+            return torch.bfloat16
+        else:
+            return self.compute_dtype
+
     @torch.no_grad()
     def _init_param_attributes(self, p: Parameter) -> None:
         """
@@ -1295,7 +1345,9 @@ def _init_param_attributes(self, p: Parameter) -> None:
             # storage to size 0 at init (here) and re-materialize (by copying
             # from _fp32_shard) as needed. If offloading params to CPU, the
             # dtype of the fp16 shard will depend on the *`compute_dtype`*.
-            p._fp16_shard = torch.zeros_like(p._fp32_shard, device=self.compute_device, dtype=self.compute_dtype)
+            p._fp16_shard = torch.zeros_like(
+                p._fp32_shard, device=self.compute_device, dtype=self._shard_dtype(p)
+            )
             free_storage_(p._fp16_shard)
 
         if self.mixed_precision:
@@ -1314,7 +1366,9 @@ def _init_param_attributes(self, p: Parameter) -> None:
         # relevant computation.
         if p._is_sharded:
             p._full_param_padded = torch.zeros(
-                p.data.numel() * self.world_size, device=self.compute_device, dtype=self.compute_dtype
+                p.data.numel() * self.world_size,
+                device=self.compute_device,
+                dtype=self._shard_dtype(p),
             )
             free_storage_(p._full_param_padded)
 
@@ -1429,7 +1483,7 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
 
         # For root and mixed precision, we convert the input to FP16 (no_grad is needed for
         # the conversion).
-        is_bf16 = self.compute_dtype == torch.bfloat16
+        is_bf16 = (self.compute_dtype == torch.bfloat16) or self._is_fp8_compute
         if self._is_root and self.mixed_precision:
             args, kwargs = cast_floats_to_right_precision(True, True, is_bf16, *args, **kwargs)
 
@@ -1443,9 +1497,13 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
         if self.force_input_to_fp32 and not self.mixed_precision:
             args, kwargs = cast_floats_to_right_precision(False, False, is_bf16, *args, **kwargs)
 
+        self.module.is_first_batch = not getattr(self, "is_not_first_batch", False)
+        self.is_not_first_batch = True
+
         # All-gather full parameters. This will also transfer FP32 parameters to
         # ``self.compute_dtype`` (e.g., FP16 if *mixed_precision* is ``True``).
         self._rebuild_full_params()
+        self.is_full_params_first_rebuilt = False
 
         if (
             self._fsdp_forward_ordering is not None
@@ -1713,8 +1771,7 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
         # then subsequent hook callbacks will see POST state.
         self.assert_state([TrainingState.BACKWARD_PRE, TrainingState.BACKWARD_POST])
         self.training_state = TrainingState.BACKWARD_POST
-        if param.grad is None:
-            return
+
 
         if hasattr(param, "_linked_param"):
             # This links to a shared param. We should finalize the linked param here.
@@ -1726,7 +1783,9 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
             if hasattr(param._linked_param, "_is_shared") and param._linked_param._is_shared:
                 param = param._linked_param
 
-        assert param.grad is not None, param.shape
+        if param.grad is None:
+            return
+
         if param.grad.requires_grad:
             raise RuntimeError("FSDP only works with gradients that don't require gradients")
 
@@ -1734,7 +1793,7 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
             # Free full params.
             self._free_full_params([param])
 
-        if self.mixed_precision and (self._require_backward_grad_sync or self.reshard_after_forward):
+        if self.mixed_precision and (self._require_backward_grad_sync or self.reshard_after_forward) and not self.fp8_allgather:
             # This is a no-op if reshard_after_forward is True, since we already
             # free the param shard when rebuilding the full params in the
             # pre_backward_hook.
@@ -1748,12 +1807,12 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                 param.unsharded_main_grad = param.grad.to(torch.float32)
             else:
                 param.unsharded_main_grad.add_(param.grad.data)
-
             param.grad = None
 
         if not self._require_backward_grad_sync:
             return
 
+
         # Wait for all work in the current stream to finish, then start the
         # reductions in post_backward stream.
         self._streams["post_backward"].wait_stream(torch.cuda.current_stream())
@@ -1800,6 +1859,19 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                 self._reducer.reduce_scatter_async(
                     grad, group=self.process_group_reduce_scatter, callback_fn=callback_fn
                 )
+            elif self.fp8_allgather:
+                if param.grad is not None:
+                    if self.fp32_reduce_scatter:
+                        param.unsharded_main_grad = param.grad.to(torch.float32)
+                    else:
+                        param.unsharded_main_grad = param.grad
+                    param.grad = None
+
+                    torch.distributed.all_reduce(
+                        param.unsharded_main_grad,
+                        group=self.process_group_reduce_scatter,
+                    )
+                    self._post_reduction_hook(param, param.unsharded_main_grad)
             else:
                 # Currently the only way for _is_sharded to be False is if
                 # world_size == 1. This could be relaxed in the future, in which
@@ -1848,8 +1920,17 @@ def _post_reduction_hook(self, param: Parameter, reduced_grad: torch.Tensor) ->
                 ), f"{param._saved_grad_shard.shape} vs {reduced_grad.shape}"
                 param._saved_grad_shard.data += reduced_grad.data
             reduced_grad = param._saved_grad_shard.data
-        elif (param.grad is None) and self.fp32_reduce_scatter:
-            param.main_grad = reduced_grad.data
+
+            if getattr(param, "unsharded_main_grad", None) is not None:
+                free_storage_(param.unsharded_main_grad.data)
+                param.unsharded_main_grad = None
+
+        elif param.grad is None:
+        # TODO(shikaili): Still having missing grad/main_grad issue.
+            if self.fp32_reduce_scatter:
+                param.main_grad = reduced_grad.data
+            else:
+                param.grad = reduced_grad.data
 
         # Optionally move gradients to CPU, typically used if one is running the optimizer on the CPU. Once the full
         # backwards pass completes, we will set `.grad` to the CPU copy.
@@ -1961,6 +2042,7 @@ def _finalize_parameters(fsdp_module: FullyShardedDataParallel) -> None:
                     # again after post-backward
                     if p.shape != p._saved_grad_shard.shape:
                         self._use_fp32_param_shard([p])
+                    assert getattr(p, "unsharded_main_grad", None) is None
                     if p._saved_grad_shard.dtype != p.dtype:
                         p.main_grad = p._saved_grad_shard
                     else:
@@ -2027,12 +2109,19 @@ def _rebuild_full_params_recursive(self):
         for module in self.modules():
             if isinstance(module, FullyShardedDataParallel):
                 module._lazy_init()
-                module._rebuild_full_params(wait_for_all_gather=False)
+                module._rebuild_full_params(
+                    wait_for_all_gather=False, not_from_recursive=False
+                )
 
 
 
     @torch.no_grad()
-    def _rebuild_full_params(self, force_full_precision: bool = False, wait_for_all_gather = True) -> Optional[List[Tuple[torch.Tensor, bool]]]:
+    def _rebuild_full_params(
+        self,
+        force_full_precision: bool = False,
+        wait_for_all_gather=True,
+        not_from_recursive=True,
+    ) -> Optional[List[Tuple[torch.Tensor, bool]]]:
         """
         Gather all shards of params.
 
@@ -2052,6 +2141,41 @@ def _rebuild_full_params(self, force_full_precision: bool = False, wait_for_all_
             caller to free the full-sized param. This will be ``None`` if
             ``force_full_precision=False`` and the full params are already gathered.
         """
+        if self._is_fp8_compute:
+            # Need to use fp32_to_fp16 stream since _cast_fp32_param_shards_to_fp16 depends on this block.
+            with torch.no_grad(), torch.cuda.stream(self._streams["fp32_to_fp16"]):
+                for p in self.params:
+                    if not isinstance(p, FlatParameter):
+                        continue
+                    d = {info[0]: info[1] for info in p._param_infos}
+                    for n, m in d.items():
+                        # Previous iteration was grad_enabled
+                        if not m.fp8_initialized:
+                            m.fp8_init(
+                                num_gemms=2 if isinstance(m, te.LayerNormMLP) else 1
+                            )
+                        if m.fp8_meta.get("update_amax_and_scale_fwd", False):
+                            if m.fp8_meta["recipe"].reduce_amax:
+                                FP8GlobalStateManager.copy_amax_from_global_buffer(
+                                    m.fp8_meta, forward=True
+                                )
+                            amax_and_scale_update(
+                                m.fp8_meta,
+                                True,
+                                update_weight_scale_inv=self.is_full_params_first_rebuilt,
+                            )
+                            if not_from_recursive:
+                                FP8GlobalStateManager.set_amax_buffer_key_deletion(
+                                    m.fp8_meta, forward=True
+                                )
+                        else:
+                            amax_and_scale_update(
+                                m.fp8_meta,
+                                True,
+                                update_weight_scale_inv=self.is_full_params_first_rebuilt,
+                            )
+                torch.cuda.current_stream().wait_stream(self._streams["fp32_to_fp16"])
+
         output_tensors: List[Tuple[torch.Tensor, bool]] = []
 
         def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
@@ -2062,6 +2186,7 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
                 custom_output_tensor (torch.Tensor, Optional): if not None, this
                 tensor contains the data we just gathered.
             """
+            p_fp16_shard_size = -1
             if custom_output_tensor is not None:
                 assert p._is_sharded
                 p.data = custom_output_tensor
@@ -2070,6 +2195,7 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
                 if (self.mixed_precision or self.move_params_to_cpu) and not force_full_precision:
                     assert p._fp16_shard is not None
                     p.data = p._fp16_shard
+                    p_fp16_shard_size = p._fp16_shard.storage().size()
                     output_tensors.append((p.data, True))
                 else:
                     # Here p.data == p._fp32_shard, so it's not safe to free.
@@ -2211,6 +2337,14 @@ def _prep_grads_for_backward(self) -> None:
         right shape, device, accumulated values, etc.
         """
         for p in self.params:
+            if isinstance(p, FlatParameter) and all(
+                _is_te_module_with_weights(info[1]) for info in p._param_infos
+            ):
+                if getattr(p, "main_grad", None) is None:
+                    p.main_grad = torch.empty_like(p, dtype=torch.float)
+                main_grad_views = p.get_param_views(p.main_grad)
+                for (_, m, n), main_grad in zip(p._param_infos, main_grad_views):
+                    getattr(m, n).main_grad = main_grad
             if p.grad is not None:
                 if p.grad.device != p.data.device:
                     p.grad = None
@@ -2232,6 +2366,8 @@ def _free_full_params(self, params: Optional[List[Parameter]] = None) -> None:
         """Free up storage for full parameters."""
         if params is None:
             params = self.params
+        self.is_not_first_batch = False
+        self.is_full_params_first_rebuilt = True
         self.has_full_params = False
 
         current_stream = torch.cuda.current_stream()
@@ -2278,8 +2414,11 @@ def local_metadata_dict(self) -> Dict[str, Any]:
                         backing_param_name = m.module.flat_param_names[i]
                         names, shapes, numels = m.module.metadata(i)
                     else:
+                        # TODO(shikaili): Understand this change.
                         assert len(m._param_name_groups[i]) == 1
-                        backing_param_name = m._param_name_groups[i][0]
+                        backing_param_name = m._param_name_groups[
+                            m._num_flatten_params
+                        ][i - m._num_flatten_params]
                         names = [backing_param_name]
                         shapes = [p._orig_size]
                         numels = [p._orig_size.numel()]
@@ -2395,12 +2534,50 @@ def _cast_fp32_param_shards_to_fp16(self, params: Optional[List[Parameter]] = No
             for p in params:
                 assert p._fp16_shard is not None
                 alloc_storage_(p._fp16_shard, size=p._fp32_shard.size())
-                p._fp16_shard.copy_(
-                    # If move_params_to_cpu is True, this will be non-blocking
-                    # because _fp32_shard is pinned, otherwise it's a no-op.
-                    p._fp32_shard.to(p._fp16_shard.device, non_blocking=True)
-                )
-                p.data = p._fp16_shard
+                if self._is_fp8_compute and _is_fp8_dtype(p._fp16_shard.dtype):
+                    assert isinstance(p, FlatParameter)
+                    assert len(p._param_infos) == len(p._param_numels)
+                    numel_per_shard = p.numel()
+                    offset = -numel_per_shard * self.rank
+                    for i in range(len(p._param_infos)):
+                        _, m, n = p._param_infos[i]
+                        numel = p._param_numels[i]
+                        if offset + numel <= 0 or offset >= numel_per_shard:
+                            offset += numel
+                            continue
+                        fp8_dtype_forward = te.fp8.get_fp8_te_dtype(
+                            m.fp8_meta["recipe"], fprop_tensor=True
+                        )
+                        if not m.fp8_initialized:
+                            m.fp8_init(
+                                num_gemms=2 if isinstance(m, te.LayerNormMLP) else 1
+                            )
+                        begin = max(offset, 0)
+                        end = min(offset + numel, numel_per_shard)
+                        cast_to_fp8(
+                            p._fp32_shard[begin:end].bfloat16(),
+                            m.fp8_meta["scaling_fwd"],
+                            (
+                                FP8FwdTensors.GEMM2_WEIGHT
+                                if n == "fc2_weight"
+                                else FP8FwdTensors.GEMM1_WEIGHT
+                            ),
+                            fp8_dtype_forward,
+                            out=p._fp16_shard[begin:end],
+                        )
+                        offset += numel
+                    p.data = p._fp16_shard.view(
+                            torch.float8_e4m3fn
+                            if fp8_dtype_forward == DType.kFloat8E4M3
+                            else torch.float8_e5m2
+                    )
+                else:
+                    p._fp16_shard.copy_(
+                        # If move_params_to_cpu is True, this will be non-blocking
+                        # because _fp32_shard is pinned, otherwise it's a no-op.
+                        p._fp32_shard.to(p._fp16_shard.device, non_blocking=True)
+                    )
+                    p.data = p._fp16_shard
         torch.cuda.current_stream().wait_stream(self._streams["fp32_to_fp16"])
 
     @torch.no_grad()
diff --git a/fairscale/nn/misc/__init__.py b/fairscale/nn/misc/__init__.py
index 71a34cae3..44999f0ca 100644
--- a/fairscale/nn/misc/__init__.py
+++ b/fairscale/nn/misc/__init__.py
@@ -9,7 +9,7 @@
 # in favor of fairscale.nn.checkpoint.checkpoint_wrapper.
 from fairscale.nn.checkpoint import checkpoint_wrapper
 
-from .flatten_params_wrapper import FlattenParamsWrapper
+from .flatten_params_wrapper import FlatParameter, FlattenParamsWrapper
 from .param_bucket import GradBucket, ParamBucket
 
 __all__: List[str] = []
diff --git a/fairscale/nn/misc/flatten_params_wrapper.py b/fairscale/nn/misc/flatten_params_wrapper.py
index 38265dd2b..ae2f3d792 100644
--- a/fairscale/nn/misc/flatten_params_wrapper.py
+++ b/fairscale/nn/misc/flatten_params_wrapper.py
@@ -486,7 +486,8 @@ def load_state_dict(
             return super().load_state_dict(state_dict, strict)
 
     def forward(self, *inputs: Any, **kwinputs: Any) -> Any:
-        self._unflatten_params_as_views()
+        if getattr(self, "is_first_batch", False):
+            self._unflatten_params_as_views()
         return self.module(*inputs, **kwinputs)
 
     def get_param_views(self, external_data_list: Optional[List[Optional[Tensor]]] = None) -> Iterator[Tensor]:

From 70f5ff540143360a5b9f67516496d33f9475f578 Mon Sep 17 00:00:00 2001
From: Shikai Li <shikaili@fb.com>
Date: Tue, 26 Mar 2024 21:28:16 -0700
Subject: [PATCH 03/20] Debug non-determinism issues.

This commit works with a 4 GPU run on SMALL model with FSDP and PP
enabled.
---
 .../fully_sharded_data_parallel.py            | 166 +++++++++---------
 fairscale/nn/misc/__init__.py                 |   2 +-
 2 files changed, 80 insertions(+), 88 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index b66c09bd8..d55e6d97c 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -46,7 +46,8 @@
 from transformer_engine.pytorch.cpp_extensions import cast_to_fp8, DType, FP8FwdTensors
 from transformer_engine.pytorch.fp8 import amax_and_scale_update, FP8GlobalStateManager
 
-from fairscale.nn.misc import FlatParameter, FlattenParamsWrapper
+from fairscale.nn.misc import FlattenParamsWrapper
+from fairscale.nn.misc.flatten_params_wrapper import FlatParameter
 from fairscale.nn.wrap import auto_wrap, config_auto_wrap_policy, enable_wrap
 from fairscale.utils.containers import apply_to_tensors
 from fairscale.utils.parallel import (
@@ -1273,7 +1274,7 @@ def _lazy_init(self) -> None:
             # ``optim.step()`` is done before we all-gather parameters.
             self._wait_for_previous_optim_step()
 
-    def _shard_dtype(self, p: Parameter) -> torch.dtype:
+    def _param_dtype(self, p: Parameter) -> torch.dtype:
         """
         Return the dtype to use for the sharded parameters.
 
@@ -1346,7 +1347,7 @@ def _init_param_attributes(self, p: Parameter) -> None:
             # from _fp32_shard) as needed. If offloading params to CPU, the
             # dtype of the fp16 shard will depend on the *`compute_dtype`*.
             p._fp16_shard = torch.zeros_like(
-                p._fp32_shard, device=self.compute_device, dtype=self._shard_dtype(p)
+                p._fp32_shard, device=self.compute_device, dtype=self._param_dtype(p)
             )
             free_storage_(p._fp16_shard)
 
@@ -1368,7 +1369,7 @@ def _init_param_attributes(self, p: Parameter) -> None:
             p._full_param_padded = torch.zeros(
                 p.data.numel() * self.world_size,
                 device=self.compute_device,
-                dtype=self._shard_dtype(p),
+                dtype=self._param_dtype(p),
             )
             free_storage_(p._full_param_padded)
 
@@ -1783,11 +1784,19 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
             if hasattr(param._linked_param, "_is_shared") and param._linked_param._is_shared:
                 param = param._linked_param
 
-        if param.grad is None:
+        # Prefer to use `param.main_grad` with higher precision in reduction than
+        # `param.grad` with equal or lower precision.
+        grad = param.grad
+        main_grad = getattr(param, "main_grad", None)
+        to_reduce_grad = main_grad if (main_grad is not None and not param.main_grad.eq(0.0).all()) else grad
+
+        if to_reduce_grad is None:
             return
 
-        if param.grad.requires_grad:
-            raise RuntimeError("FSDP only works with gradients that don't require gradients")
+        if to_reduce_grad.requires_grad:
+            raise RuntimeError(
+                    "FSDP only works with gradients that don't require gradients"
+                    )
 
         if self._should_free_in_backward():
             # Free full params.
@@ -1802,49 +1811,52 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
         # Switch to FP32 shard after backward.
         self._use_fp32_param_shard([param])
 
+        # Accumulate gradients manually if in FP32 instead of default precision.
         if self.fp32_reduce_scatter:
-            if getattr(param, "unsharded_main_grad", None) is None:
-                param.unsharded_main_grad = param.grad.to(torch.float32)
-            else:
-                param.unsharded_main_grad.add_(param.grad.data)
-            param.grad = None
+            # If `param.grad` is `None`, then fp32 reduction is already happening at
+            # `param.main_grad`, nothing need to be done here.
+            if param.grad is not None:
+                if getattr(param, "unsharded_main_grad", None) is None:
+                    param.unsharded_main_grad = param.grad.to(torch.float32)
+                else:
+                    param.unsharded_main_grad.add_(param.grad.to(torch.float32))
+                # Resets `param.grad` to avoid PyTorch accumulation.
+                param.grad = None
 
         if not self._require_backward_grad_sync:
             return
 
-
         # Wait for all work in the current stream to finish, then start the
         # reductions in post_backward stream.
         self._streams["post_backward"].wait_stream(torch.cuda.current_stream())
         with torch.cuda.stream(self._streams["post_backward"]):
 
-            if self.fp32_reduce_scatter:
-                # Cast grad to FP32.
-                orig_grad_data = param.unsharded_main_grad.data
-            else:
-                orig_grad_data = param.grad.data
+            # Prefer to use "unsharded_main_grad" with higher precision if exsits.
+            unsharded_main_grad = getattr(param, "unsharded_main_grad", None)
+            to_reduce_grad = unsharded_main_grad if unsharded_main_grad is not None else to_reduce_grad
+
+            # Clear grad on the tensor, so any repeated gradient computations do not interfere with this reduction.
+            # 1. For sharded parameters, we will asynchronously accumulate the reduced gradient into
+            # `param._saved_grad_shard` which will be re-sharded to `param.grad`/`param.main_grad` later after
+            # finalization.
+            # 2. For unsharded parameters, we will directly asynchronously accumulate the reduce gradient into
+            # `param.grad`/`param.main_grad`.
+            param.grad = None
+            param.main_grad = None
+            # `param.unsharded_main_grad` is no longer usefual and will be created again inside
+            # `_post_backward_hook` in the next first microbatch.
+            param.unsharded_main_grad = None
 
             if self.gradient_predivide_factor > 1:
                 # Average grad by world_size for consistency with PyTorch DDP.
-                if getattr(param, "unsharded_main_grad", None) is not None:
-                    param.unsharded_main_grad.data.div_(self.gradient_predivide_factor)
-                else:
-                    param.grad.data.div_(self.gradient_predivide_factor)
+                to_reduce_grad.data.div_(self.gradient_predivide_factor)
 
             if param._is_sharded:
                 assert self._reducer is not None
-                # Save the unsharded grad for reduction. We will asynchronously accumulate the reduced gradient into
-                # param._saved_grad_shard. If this FSDP module was called multiple times it's possible that multiple
-                # gradient reductions will happen in an undefined order. But addition commutes, so this order doesn't
-                # matter, neglecting rounding.
-                if getattr(param, "unsharded_main_grad", None) is not None:
-                    grad = param.unsharded_main_grad.data
-                    param.unsharded_main_grad = None
-                else:
-                    grad = param.grad.data
-                    param.grad = None
 
-                # Clear grad on the tensor, so any repeated gradient computations do not interfere with this reduction.
+                # If this FSDP module was called multiple times it's possible that multiple gradient reductions will
+                # happen in an undefined order. But addition commutes, so this order doesn't matter, neglecting
+                # rounding.
                 #
                 # The effect on memory consumption is not usually significant. No extra memory is allocated if this
                 # module is called only once, reduction happens quickly, or the tensor is bucketed. If the module is
@@ -1857,37 +1869,29 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                 # unsharded gradients allocated; one for a pending reduction, and one for gradient computation.
                 callback_fn = functools.partial(self._post_reduction_hook, param)
                 self._reducer.reduce_scatter_async(
-                    grad, group=self.process_group_reduce_scatter, callback_fn=callback_fn
+                    to_reduce_grad,
+                    group=self.process_group_reduce_scatter,
+                    callback_fn=callback_fn,
                 )
-            elif self.fp8_allgather:
-                if param.grad is not None:
-                    if self.fp32_reduce_scatter:
-                        param.unsharded_main_grad = param.grad.to(torch.float32)
-                    else:
-                        param.unsharded_main_grad = param.grad
-                    param.grad = None
-
+            else:
+                # Unsharded parameters only happens for word_size == 1 or fp8_allgather
+                assert self.world_size == 1 or (self._is_fp8_compute and not isinstance(param, FlatParameter))
+                if self.world_size > 1:
                     torch.distributed.all_reduce(
-                        param.unsharded_main_grad,
+                        to_reduce_grad,
                         group=self.process_group_reduce_scatter,
                     )
-                    self._post_reduction_hook(param, param.unsharded_main_grad)
-            else:
-                # Currently the only way for _is_sharded to be False is if
-                # world_size == 1. This could be relaxed in the future, in which
-                # case grads should be all-reduced here.
-                assert self.world_size == 1
-                if getattr(param, "unsharded_main_grad", None) is not None:
-                    self._post_reduction_hook(param, param.unsharded_main_grad)
-                else:
-                    self._post_reduction_hook(param, param.grad)
+                self._post_reduction_hook(param, to_reduce_grad)
 
             # After _post_backward_hook returns, orig_grad_data will eventually
             # go out of scope, at which point it could otherwise be freed for
             # further reuse by the main stream while the div/reduce_scatter/copy
             # are underway in the post_backward stream. See:
             # github.com/NVIDIA/apex/blob/master/apex/parallel/distributed.py
-            orig_grad_data.record_stream(self._streams["post_backward"])
+            for g in (grad, main_grad, unsharded_main_grad):
+                if g is not None:
+                    g.data.record_stream(self._streams["post_backward"])
+
             if self._reduce_scatter_free_event_queue is not None:
                 release_full_grad_event = torch.cuda.Event()
                 release_full_grad_event.record()
@@ -1898,17 +1902,12 @@ def _post_reduction_hook(self, param: Parameter, reduced_grad: torch.Tensor) ->
         """Hook to call on each param after the reduce-scatter."""
         assert torch.cuda.current_stream() == self._streams["post_backward"]
         self.assert_state(TrainingState.BACKWARD_POST)
+
+        assert not (self.fp32_reduce_scatter and reduced_grad.dtype != param.dtype)
+
         if self.gradient_postdivide_factor > 1:
             # Average grad by world_size for consistency with PyTorch DDP.
             reduced_grad.data.div_(self.gradient_postdivide_factor)
-        # Cast grad to param's dtype (typically FP32). Note: we do this
-        # before the move_grads_to_cpu step so that this entire hook remains
-        # non-blocking. The downside is a bit more D2H transfer in that case.
-        if self.fp32_reduce_scatter:
-            orig_param_grad_data = reduced_grad.data
-            # reduced_grad.data = reduced_grad.data.to(dtype=param.data.dtype)
-            # Don't let this memory get reused until after the transfer.
-            orig_param_grad_data.record_stream(torch.cuda.current_stream())
 
         if param._is_sharded:
             # Accumulate into the gradient shard.
@@ -1921,16 +1920,9 @@ def _post_reduction_hook(self, param: Parameter, reduced_grad: torch.Tensor) ->
                 param._saved_grad_shard.data += reduced_grad.data
             reduced_grad = param._saved_grad_shard.data
 
-            if getattr(param, "unsharded_main_grad", None) is not None:
-                free_storage_(param.unsharded_main_grad.data)
-                param.unsharded_main_grad = None
-
         elif param.grad is None:
-        # TODO(shikaili): Still having missing grad/main_grad issue.
             if self.fp32_reduce_scatter:
                 param.main_grad = reduced_grad.data
-            else:
-                param.grad = reduced_grad.data
 
         # Optionally move gradients to CPU, typically used if one is running the optimizer on the CPU. Once the full
         # backwards pass completes, we will set `.grad` to the CPU copy.
@@ -2150,6 +2142,7 @@ def _rebuild_full_params(
                     d = {info[0]: info[1] for info in p._param_infos}
                     for n, m in d.items():
                         # Previous iteration was grad_enabled
+                        # assert m.fp8_initialized
                         if not m.fp8_initialized:
                             m.fp8_init(
                                 num_gemms=2 if isinstance(m, te.LayerNormMLP) else 1
@@ -2159,22 +2152,22 @@ def _rebuild_full_params(
                                 FP8GlobalStateManager.copy_amax_from_global_buffer(
                                     m.fp8_meta, forward=True
                                 )
-                            amax_and_scale_update(
-                                m.fp8_meta,
-                                True,
-                                update_weight_scale_inv=self.is_full_params_first_rebuilt,
-                            )
-                            if not_from_recursive:
-                                FP8GlobalStateManager.set_amax_buffer_key_deletion(
-                                    m.fp8_meta, forward=True
+                                amax_and_scale_update(
+                                    m.fp8_meta,
+                                    True,
+                                    update_weight_scale_inv=self.is_full_params_first_rebuilt,
                                 )
-                        else:
-                            amax_and_scale_update(
-                                m.fp8_meta,
-                                True,
-                                update_weight_scale_inv=self.is_full_params_first_rebuilt,
-                            )
-                torch.cuda.current_stream().wait_stream(self._streams["fp32_to_fp16"])
+                                if not_from_recursive:
+                                    FP8GlobalStateManager.set_amax_buffer_key_deletion(
+                                        m.fp8_meta, forward=True
+                                    )
+                            else:
+                                amax_and_scale_update(
+                                    m.fp8_meta,
+                                    True,
+                                    update_weight_scale_inv=self.is_full_params_first_rebuilt,
+                                )
+            torch.cuda.current_stream().wait_stream(self._streams["fp32_to_fp16"])
 
         output_tensors: List[Tuple[torch.Tensor, bool]] = []
 
@@ -2341,7 +2334,7 @@ def _prep_grads_for_backward(self) -> None:
                 _is_te_module_with_weights(info[1]) for info in p._param_infos
             ):
                 if getattr(p, "main_grad", None) is None:
-                    p.main_grad = torch.empty_like(p, dtype=torch.float)
+                    p.main_grad = torch.empty_like(p, dtype=torch.float32)
                 main_grad_views = p.get_param_views(p.main_grad)
                 for (_, m, n), main_grad in zip(p._param_infos, main_grad_views):
                     getattr(m, n).main_grad = main_grad
@@ -2414,8 +2407,6 @@ def local_metadata_dict(self) -> Dict[str, Any]:
                         backing_param_name = m.module.flat_param_names[i]
                         names, shapes, numels = m.module.metadata(i)
                     else:
-                        # TODO(shikaili): Understand this change.
-                        assert len(m._param_name_groups[i]) == 1
                         backing_param_name = m._param_name_groups[
                             m._num_flatten_params
                         ][i - m._num_flatten_params]
@@ -2545,6 +2536,7 @@ def _cast_fp32_param_shards_to_fp16(self, params: Optional[List[Parameter]] = No
                         if offset + numel <= 0 or offset >= numel_per_shard:
                             offset += numel
                             continue
+                        assert _is_te_module_with_weights(m)
                         fp8_dtype_forward = te.fp8.get_fp8_te_dtype(
                             m.fp8_meta["recipe"], fprop_tensor=True
                         )
diff --git a/fairscale/nn/misc/__init__.py b/fairscale/nn/misc/__init__.py
index 44999f0ca..71a34cae3 100644
--- a/fairscale/nn/misc/__init__.py
+++ b/fairscale/nn/misc/__init__.py
@@ -9,7 +9,7 @@
 # in favor of fairscale.nn.checkpoint.checkpoint_wrapper.
 from fairscale.nn.checkpoint import checkpoint_wrapper
 
-from .flatten_params_wrapper import FlatParameter, FlattenParamsWrapper
+from .flatten_params_wrapper import FlattenParamsWrapper
 from .param_bucket import GradBucket, ParamBucket
 
 __all__: List[str] = []

From 16c682db1474aea4fe85a2be9d1a87b3bc6434f0 Mon Sep 17 00:00:00 2001
From: Shikai Li <shikaili@fb.com>
Date: Thu, 28 Mar 2024 16:22:18 -0700
Subject: [PATCH 04/20] Moves amax update logic into params downcasting
 function.

---
 .../fully_sharded_data_parallel.py            | 92 +++++++++----------
 1 file changed, 44 insertions(+), 48 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index d55e6d97c..99b3b6234 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -2133,42 +2133,6 @@ def _rebuild_full_params(
             caller to free the full-sized param. This will be ``None`` if
             ``force_full_precision=False`` and the full params are already gathered.
         """
-        if self._is_fp8_compute:
-            # Need to use fp32_to_fp16 stream since _cast_fp32_param_shards_to_fp16 depends on this block.
-            with torch.no_grad(), torch.cuda.stream(self._streams["fp32_to_fp16"]):
-                for p in self.params:
-                    if not isinstance(p, FlatParameter):
-                        continue
-                    d = {info[0]: info[1] for info in p._param_infos}
-                    for n, m in d.items():
-                        # Previous iteration was grad_enabled
-                        # assert m.fp8_initialized
-                        if not m.fp8_initialized:
-                            m.fp8_init(
-                                num_gemms=2 if isinstance(m, te.LayerNormMLP) else 1
-                            )
-                        if m.fp8_meta.get("update_amax_and_scale_fwd", False):
-                            if m.fp8_meta["recipe"].reduce_amax:
-                                FP8GlobalStateManager.copy_amax_from_global_buffer(
-                                    m.fp8_meta, forward=True
-                                )
-                                amax_and_scale_update(
-                                    m.fp8_meta,
-                                    True,
-                                    update_weight_scale_inv=self.is_full_params_first_rebuilt,
-                                )
-                                if not_from_recursive:
-                                    FP8GlobalStateManager.set_amax_buffer_key_deletion(
-                                        m.fp8_meta, forward=True
-                                    )
-                            else:
-                                amax_and_scale_update(
-                                    m.fp8_meta,
-                                    True,
-                                    update_weight_scale_inv=self.is_full_params_first_rebuilt,
-                                )
-            torch.cuda.current_stream().wait_stream(self._streams["fp32_to_fp16"])
-
         output_tensors: List[Tuple[torch.Tensor, bool]] = []
 
         def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
@@ -2179,7 +2143,6 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
                 custom_output_tensor (torch.Tensor, Optional): if not None, this
                 tensor contains the data we just gathered.
             """
-            p_fp16_shard_size = -1
             if custom_output_tensor is not None:
                 assert p._is_sharded
                 p.data = custom_output_tensor
@@ -2188,7 +2151,6 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
                 if (self.mixed_precision or self.move_params_to_cpu) and not force_full_precision:
                     assert p._fp16_shard is not None
                     p.data = p._fp16_shard
-                    p_fp16_shard_size = p._fp16_shard.storage().size()
                     output_tensors.append((p.data, True))
                 else:
                     # Here p.data == p._fp32_shard, so it's not safe to free.
@@ -2237,8 +2199,10 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
                 event.synchronize()
 
         with torch.cuda.stream(self._streams["all_gather"]):
-            if (self.mixed_precision or self.move_params_to_cpu) and not force_full_precision:
-                self._cast_fp32_param_shards_to_fp16()
+            if (
+                self.mixed_precision or self.move_params_to_cpu
+            ) and not force_full_precision:
+                self._cast_fp32_param_shards_to_fp16(not_from_recursive=not_from_recursive)
 
             if self.move_params_to_cpu:
                 if force_full_precision:
@@ -2246,7 +2210,7 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
                     # use pinned memory. Otherwise move p.data to the compute
                     # device.
                     if self.params[0].dtype == self.compute_dtype:
-                        self._cast_fp32_param_shards_to_fp16()
+                        self._cast_fp32_param_shards_to_fp16(not_from_recursive=not_from_recursive)
                     else:
                         for p in self.params:
                             p.data = p.data.to(self.compute_device)
@@ -2517,33 +2481,64 @@ def _use_fp32_param_shard(self, params: Optional[List[Parameter]] = None) -> Non
             p.data = p._fp32_shard
 
     @torch.no_grad()
-    def _cast_fp32_param_shards_to_fp16(self, params: Optional[List[Parameter]] = None) -> None:
+    def _cast_fp32_param_shards_to_fp16(
+        self, params: Optional[List[Parameter]] = None,
+        not_from_recursive: bool = False,
+    ) -> None:
         """Cast FP32 param shard to FP16 for a list of params."""
         if params is None:
             params = self.params
+
         with torch.cuda.stream(self._streams["fp32_to_fp16"]):
             for p in params:
                 assert p._fp16_shard is not None
                 alloc_storage_(p._fp16_shard, size=p._fp32_shard.size())
+
                 if self._is_fp8_compute and _is_fp8_dtype(p._fp16_shard.dtype):
                     assert isinstance(p, FlatParameter)
                     assert len(p._param_infos) == len(p._param_numels)
+
                     numel_per_shard = p.numel()
                     offset = -numel_per_shard * self.rank
                     for i in range(len(p._param_infos)):
                         _, m, n = p._param_infos[i]
+                        assert _is_te_module_with_weights(m)
+
+                        if not m.fp8_initialized:
+                            m.fp8_init(
+                                num_gemms=2 if isinstance(m, te.LayerNormMLP) else 1
+                            )
+
+                        if self.is_full_params_first_rebuilt:
+                            if m.fp8_meta.get("update_amax_and_scale_fwd", False):
+                                if m.fp8_meta["recipe"].reduce_amax:
+                                    FP8GlobalStateManager.copy_amax_from_global_buffer(
+                                        m.fp8_meta, forward=True
+                                    )
+                                    amax_and_scale_update(
+                                        m.fp8_meta,
+                                        True,
+                                        update_weight_scale_inv=True,
+                                    )
+                                    if not_from_recursive:
+                                        FP8GlobalStateManager.set_amax_buffer_key_deletion(
+                                            m.fp8_meta, forward=True
+                                        )
+                                else:
+                                    amax_and_scale_update(
+                                        m.fp8_meta,
+                                        True,
+                                        update_weight_scale_inv=True,
+                                    )
+
                         numel = p._param_numels[i]
                         if offset + numel <= 0 or offset >= numel_per_shard:
                             offset += numel
                             continue
-                        assert _is_te_module_with_weights(m)
+
                         fp8_dtype_forward = te.fp8.get_fp8_te_dtype(
                             m.fp8_meta["recipe"], fprop_tensor=True
                         )
-                        if not m.fp8_initialized:
-                            m.fp8_init(
-                                num_gemms=2 if isinstance(m, te.LayerNormMLP) else 1
-                            )
                         begin = max(offset, 0)
                         end = min(offset + numel, numel_per_shard)
                         cast_to_fp8(
@@ -2570,7 +2565,8 @@ def _cast_fp32_param_shards_to_fp16(self, params: Optional[List[Parameter]] = No
                         p._fp32_shard.to(p._fp16_shard.device, non_blocking=True)
                     )
                     p.data = p._fp16_shard
-        torch.cuda.current_stream().wait_stream(self._streams["fp32_to_fp16"])
+
+        self._streams["all_gather"].wait_stream(self._streams["fp32_to_fp16"])
 
     @torch.no_grad()
     def _free_fp16_param_shard(self, params: Optional[List[Parameter]] = None) -> None:

From 24a769ff1ba0014aa2c4f32c7c749bfefb954ab1 Mon Sep 17 00:00:00 2001
From: Shikai Li <shikaili@fb.com>
Date: Thu, 28 Mar 2024 17:13:57 -0700
Subject: [PATCH 05/20] Cleans up code.

---
 .../fully_sharded_data_parallel.py            | 67 +++++++++----------
 1 file changed, 33 insertions(+), 34 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index 99b3b6234..354022499 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -267,6 +267,9 @@ class FullyShardedDataParallel(nn.Module):
         fp32_reduce_scatter (bool, Optional):
             if ``True``, then reduce-scatter gradients in FP32. This is only
             relevant when *``mixed_precision``* is ``True``.
+        fp8_all_gather (bool, Optional):
+            if ``True``, then all-gather weights/gradients in FP8. This is only
+            relevant when *``mixed_precision``* is ``True``.
         flatten_parameters (bool, Optional):
             if ``True``, flatten parameters into a single contiguous tensor,
             which improves training speed.
@@ -362,6 +365,7 @@ def __init__(
         disable_reshard_on_root: bool = True,
         mixed_precision: bool = False,
         fp32_reduce_scatter: bool = False,
+        fp8_all_gather: bool = False,
         flatten_parameters: bool = True,
         move_params_to_cpu: bool = False,
         compute_dtype: Optional[torch.dtype] = None,
@@ -381,7 +385,6 @@ def __init__(
         limit_all_gather_events: bool = False,
         limit_reduce_scatter_events: bool = False,
         should_validate_process_group: bool = True,
-        fp8_allgather: bool = False,
     ):
         try:
             import torch._C
@@ -433,6 +436,7 @@ def __init__(
         self.disable_reshard_on_root = disable_reshard_on_root
         self.mixed_precision = mixed_precision
         self.fp32_reduce_scatter = fp32_reduce_scatter
+        self.fp8_all_gather = fp8_all_gather
         self.flatten_parameters = flatten_parameters
         self.move_params_to_cpu = move_params_to_cpu or cpu_offload
         self.compute_dtype = compute_dtype or (torch.float16 if mixed_precision else torch.float32)
@@ -447,7 +451,6 @@ def __init__(
         self.force_input_to_fp32 = force_input_to_fp32
         self.verbose = verbose
         self.state_dict_on_rank_0_only = state_dict_on_rank_0_only
-        self.fp8_allgather = fp8_allgather
         # Experimental feature for now. Use at your own risk.
         self.ssd_offload = True if offload_config and offload_config.offload_type == "ssd_offload" else False
 
@@ -585,6 +588,7 @@ def __init__(
         self.dont_wait_current_stream_for_post_all_gather = False
         self._all_gather_free_event_queue = _FreeEventQueue() if limit_all_gather_events else None
         self._reduce_scatter_free_event_queue = _FreeEventQueue() if limit_reduce_scatter_events else None
+        self.is_full_params_first_rebuilt = True
 
     @property
     def _is_fp8_compute(self) -> bool:
@@ -1281,9 +1285,8 @@ def _param_dtype(self, p: Parameter) -> torch.dtype:
         Returns:
             The dtype to use for the sharded parameters.
         """
-        if self._is_fp8_compute and not isinstance(p, FlatParameter):
-            # Assume non flattened are precision critical like norm
-            assert not p._is_sharded
+        if self._is_fp8_compute and (not self.fp8_all_gather or
+                                     not isinstance(p, FlatParameter)):
             return torch.bfloat16
         else:
             return self.compute_dtype
@@ -1300,13 +1303,13 @@ def _init_param_attributes(self, p: Parameter) -> None:
             ``_orig_size``: the size of the original Parameter (before sharding)
 
         The remaining attributes are set here:
-            ``_fp32_shard``: a single shard of the parameters in full precision
-                (typically FP32, but this is dependent on the dtype of the model
-                as it's passed in by the user). This can be on CPU or GPU
+            ``_fp32_shard``: This will be a single shard of the parameters in
+                full precision (typically FP32, but this is dependent on the dtype of
+                the model as it's passed in by the user). This can be on CPU or GPU
                 depending on the value of *``move_params_to_cpu``*.
-            ``_fp16_shard``: This will be a single shard of the parameters in FP16, used for all-gather.
-                This can be in FP16 or FP32 depending on the value of *``compute_dtype``* and
-                if params are offloaded to CPU.
+            ``_fp16_shard``: This will be a single shard of the parameters
+                used for all-gather. This can be in FP8, FP16 or FP32 depending on the value
+                of *``compute_dtype``*, *``fp8_all_gather``*, *``move_params_to_cpu``*..
             ``_full_param_padded``: the full weight (padded to be evenly
                 divisible by ``world_size``), used for computation in the
                 forward and backward pass. This will be resized in place and
@@ -1326,6 +1329,7 @@ def _init_param_attributes(self, p: Parameter) -> None:
 
         if self.mixed_precision:
             assert p._fp32_shard.dtype == torch.float32
+
         if self.move_params_to_cpu:
             assert p._fp32_shard.device == torch.device("cpu")
 
@@ -1339,7 +1343,6 @@ def _init_param_attributes(self, p: Parameter) -> None:
                 p.data = p._fp32_shard
 
         if self.move_params_to_cpu or self.mixed_precision:
-
             # In mixed precision mode, we maintain a reduced precision
             # (typically FP16) parameter shard on compute_device for performing
             # the computation in the forward/backward pass. We resize the
@@ -1350,11 +1353,7 @@ def _init_param_attributes(self, p: Parameter) -> None:
                 p._fp32_shard, device=self.compute_device, dtype=self._param_dtype(p)
             )
             free_storage_(p._fp16_shard)
-
-        if self.mixed_precision:
-            assert p._fp32_shard.dtype == torch.float32
-
-        if not self.mixed_precision and not self.move_params_to_cpu:
+        else:
             # use _fp32_shard if you are not in using mixed precision or
             # offloading params and grads to CPU.
             p._fp16_shard = None
@@ -1435,8 +1434,8 @@ def _setup_streams(self) -> None:
             return
 
         if torch.cuda.is_available():
-            # Stream to move main FP32 params (may be on CPU) to FP16 for forward.
-            self._streams["fp32_to_fp16"] = torch.cuda.Stream()
+            # Stream to move main FP32 params (may be on CPU) to FP32/FP16/FP8 for forward.
+            self._streams["cast_param"] = torch.cuda.Stream()
             # Stream for all-gathering parameters.
             self._streams["all_gather"] = torch.cuda.Stream()
             # Stream for overlapping grad reduction with the backward pass.
@@ -1472,7 +1471,7 @@ def _wait_for_previous_optim_step(self) -> None:
         if not torch.cuda.is_available():
             return
         if self.mixed_precision or self.move_params_to_cpu:
-            self._streams["fp32_to_fp16"].wait_stream(torch.cuda.current_stream())
+            self._streams["cast_param"].wait_stream(torch.cuda.current_stream())
         else:
             self._streams["all_gather"].wait_stream(torch.cuda.current_stream())
 
@@ -1802,7 +1801,7 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
             # Free full params.
             self._free_full_params([param])
 
-        if self.mixed_precision and (self._require_backward_grad_sync or self.reshard_after_forward) and not self.fp8_allgather:
+        if self.mixed_precision and (self._require_backward_grad_sync or self.reshard_after_forward) and not self.fp8_all_gather:
             # This is a no-op if reshard_after_forward is True, since we already
             # free the param shard when rebuilding the full params in the
             # pre_backward_hook.
@@ -1819,7 +1818,7 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                 if getattr(param, "unsharded_main_grad", None) is None:
                     param.unsharded_main_grad = param.grad.to(torch.float32)
                 else:
-                    param.unsharded_main_grad.add_(param.grad.to(torch.float32))
+                    param.unsharded_main_grad.data.add_(param.grad.to(torch.float32))
                 # Resets `param.grad` to avoid PyTorch accumulation.
                 param.grad = None
 
@@ -1874,7 +1873,7 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                     callback_fn=callback_fn,
                 )
             else:
-                # Unsharded parameters only happens for word_size == 1 or fp8_allgather
+                # Unsharded parameters only happens for word_size == 1 or fp8_all_gather
                 assert self.world_size == 1 or (self._is_fp8_compute and not isinstance(param, FlatParameter))
                 if self.world_size > 1:
                     torch.distributed.all_reduce(
@@ -1888,9 +1887,7 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
             # further reuse by the main stream while the div/reduce_scatter/copy
             # are underway in the post_backward stream. See:
             # github.com/NVIDIA/apex/blob/master/apex/parallel/distributed.py
-            for g in (grad, main_grad, unsharded_main_grad):
-                if g is not None:
-                    g.data.record_stream(self._streams["post_backward"])
+            to_reduce_grad.data.record_stream(self._streams["post_backward"])
 
             if self._reduce_scatter_free_event_queue is not None:
                 release_full_grad_event = torch.cuda.Event()
@@ -1903,7 +1900,7 @@ def _post_reduction_hook(self, param: Parameter, reduced_grad: torch.Tensor) ->
         assert torch.cuda.current_stream() == self._streams["post_backward"]
         self.assert_state(TrainingState.BACKWARD_POST)
 
-        assert not (self.fp32_reduce_scatter and reduced_grad.dtype != param.dtype)
+        # assert not (self.fp32_reduce_scatter and reduced_grad.dtype != param.dtype)
 
         if self.gradient_postdivide_factor > 1:
             # Average grad by world_size for consistency with PyTorch DDP.
@@ -1917,12 +1914,14 @@ def _post_reduction_hook(self, param: Parameter, reduced_grad: torch.Tensor) ->
                 assert (
                     param._saved_grad_shard.shape == reduced_grad.shape
                 ), f"{param._saved_grad_shard.shape} vs {reduced_grad.shape}"
-                param._saved_grad_shard.data += reduced_grad.data
+                param._saved_grad_shard.data.add_(reduced_grad.data)
             reduced_grad = param._saved_grad_shard.data
 
         elif param.grad is None:
             if self.fp32_reduce_scatter:
                 param.main_grad = reduced_grad.data
+            else:
+                param.grad = reduced_grad.data
 
         # Optionally move gradients to CPU, typically used if one is running the optimizer on the CPU. Once the full
         # backwards pass completes, we will set `.grad` to the CPU copy.
@@ -2202,7 +2201,7 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
             if (
                 self.mixed_precision or self.move_params_to_cpu
             ) and not force_full_precision:
-                self._cast_fp32_param_shards_to_fp16(not_from_recursive=not_from_recursive)
+                self._cast_params_for_all_gather(not_from_recursive=not_from_recursive)
 
             if self.move_params_to_cpu:
                 if force_full_precision:
@@ -2210,7 +2209,7 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
                     # use pinned memory. Otherwise move p.data to the compute
                     # device.
                     if self.params[0].dtype == self.compute_dtype:
-                        self._cast_fp32_param_shards_to_fp16(not_from_recursive=not_from_recursive)
+                        self._cast_params_for_all_gather(not_from_recursive=not_from_recursive)
                     else:
                         for p in self.params:
                             p.data = p.data.to(self.compute_device)
@@ -2481,7 +2480,7 @@ def _use_fp32_param_shard(self, params: Optional[List[Parameter]] = None) -> Non
             p.data = p._fp32_shard
 
     @torch.no_grad()
-    def _cast_fp32_param_shards_to_fp16(
+    def _cast_params_for_all_gather(
         self, params: Optional[List[Parameter]] = None,
         not_from_recursive: bool = False,
     ) -> None:
@@ -2489,7 +2488,7 @@ def _cast_fp32_param_shards_to_fp16(
         if params is None:
             params = self.params
 
-        with torch.cuda.stream(self._streams["fp32_to_fp16"]):
+        with torch.cuda.stream(self._streams["cast_param"]):
             for p in params:
                 assert p._fp16_shard is not None
                 alloc_storage_(p._fp16_shard, size=p._fp32_shard.size())
@@ -2566,7 +2565,7 @@ def _cast_fp32_param_shards_to_fp16(
                     )
                     p.data = p._fp16_shard
 
-        self._streams["all_gather"].wait_stream(self._streams["fp32_to_fp16"])
+        self._streams["all_gather"].wait_stream(self._streams["cast_param"])
 
     @torch.no_grad()
     def _free_fp16_param_shard(self, params: Optional[List[Parameter]] = None) -> None:
@@ -2576,7 +2575,7 @@ def _free_fp16_param_shard(self, params: Optional[List[Parameter]] = None) -> No
         current_stream = torch.cuda.current_stream()
         for p in params:
             if p._fp16_shard is not None:
-                # _fp16_shard is allocated in "fp32_to_fp16" stream, so we can't
+                # _fp16_shard is allocated in "cast_param" stream, so we can't
                 # free it until the work in the current stream completes.
                 p._fp16_shard.record_stream(current_stream)
                 free_storage_(p._fp16_shard)

From 3e2e77f05a3d8c11a993d94cdb8008eaa9cf2313 Mon Sep 17 00:00:00 2001
From: Shikai Li <shikaili@fb.com>
Date: Mon, 1 Apr 2024 21:32:45 -0700
Subject: [PATCH 06/20] Fix `main_grad` attribute checking.

- Clean up flatten and non_flatten parameter generation logic.
- Avoid checking `main_grad` attribute all equal to zeros.
---
 .../fully_sharded_data_parallel.py            | 65 ++++++++++---------
 1 file changed, 33 insertions(+), 32 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index 354022499..f3fa2acb8 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -501,26 +501,27 @@ def __init__(
 
         # For now, it is either all flatten or none flatten. This will be extended to
         # multiple flatten groups in my next PR.
-        to_be_flatten_params: List[List[Parameter]] = [[]]
-        non_flatten_params = params
-        param_name_groups = [[n] for n in param_names]
-        if self.flatten_parameters:
-            to_be_flatten_params = [
-                [
-                    params[i]
-                    for i in range(len(params))
-                    if "norm_weight" not in param_names[i]
-                ]
-            ]
-            non_flatten_params = [
-                params[i]
-                for i in range(len(params))
-                if "norm_weight" in param_names[i]
-            ]
-            param_name_groups = [
-                [n for n in param_names if "norm_weight" not in n],
-                [n for n in param_names if "norm_weight" in n],
-            ]
+        def should_flatten(name: str) -> bool:
+            # `*_norm_weights` are numerics-sensitive and cannot be quantized to fp8.
+            return self.flatten_parameters and (not self.fp8_all_gather or "norm_weight" not in name)
+
+        to_be_flatten_params: List[List[Parameter]] = [
+            param
+            for param, name in zip(params, param_names)
+            if should_flatten(name)
+        ]
+        if to_be_flatten_params:
+            to_be_flatten_params = [to_be_flatten_params]
+        non_flatten_params: List[List[Parameter]] = [
+            param
+            for param, name in zip(params, param_names)
+            if not should_flatten(name)
+        ]
+        param_name_groups: List[List[str]] = [
+            [n for n in param_names if should_flatten(n)]
+        ] + [
+            [n] for n in param_names if not should_flatten(n)
+        ]
         del param_names
 
         self._fsdp_wrapped_module: nn.Module = FlattenParamsWrapper(
@@ -1772,7 +1773,6 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
         self.assert_state([TrainingState.BACKWARD_PRE, TrainingState.BACKWARD_POST])
         self.training_state = TrainingState.BACKWARD_POST
 
-
         if hasattr(param, "_linked_param"):
             # This links to a shared param. We should finalize the linked param here.
             assert param.shape == (1,), param.shape
@@ -1783,11 +1783,13 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
             if hasattr(param._linked_param, "_is_shared") and param._linked_param._is_shared:
                 param = param._linked_param
 
-        # Prefer to use `param.main_grad` with higher precision in reduction than
-        # `param.grad` with equal or lower precision.
         grad = param.grad
         main_grad = getattr(param, "main_grad", None)
-        to_reduce_grad = main_grad if (main_grad is not None and not param.main_grad.eq(0.0).all()) else grad
+        # Only one of `grad` or `main_grad` can exists. Whenever `main_grad is used for accumulation,
+        # grad should be set as `None`.
+        assert not (grad is not None and main_grad is not None)
+        # Use `grad` or `main_grad` that is not None and avoid invoking a kernel to check all zeros.
+        to_reduce_grad = grad if grad is not None else main_grad
 
         if to_reduce_grad is None:
             return
@@ -1873,8 +1875,8 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                     callback_fn=callback_fn,
                 )
             else:
-                # Unsharded parameters only happens for word_size == 1 or fp8_all_gather
-                assert self.world_size == 1 or (self._is_fp8_compute and not isinstance(param, FlatParameter))
+                # Unsharded parameters only happens for word_size == 1 or self.fp8_all_gather
+                assert self.world_size == 1 or (self._is_fp8_compute and self.fp8_all_gather and not isinstance(param, FlatParameter))
                 if self.world_size > 1:
                     torch.distributed.all_reduce(
                         to_reduce_grad,
@@ -2293,9 +2295,10 @@ def _prep_grads_for_backward(self) -> None:
         right shape, device, accumulated values, etc.
         """
         for p in self.params:
-            if isinstance(p, FlatParameter) and all(
-                _is_te_module_with_weights(info[1]) for info in p._param_infos
-            ):
+            fused_wgard_accumulation = (self.fp8_all_gather
+                and isinstance(p, FlatParameter)
+                and all(_is_te_module_with_weights(info[1]) for info in p._param_infos))
+            if fused_wgard_accumulation:
                 if getattr(p, "main_grad", None) is None:
                     p.main_grad = torch.empty_like(p, dtype=torch.float32)
                 main_grad_views = p.get_param_views(p.main_grad)
@@ -2370,9 +2373,7 @@ def local_metadata_dict(self) -> Dict[str, Any]:
                         backing_param_name = m.module.flat_param_names[i]
                         names, shapes, numels = m.module.metadata(i)
                     else:
-                        backing_param_name = m._param_name_groups[
-                            m._num_flatten_params
-                        ][i - m._num_flatten_params]
+                        backing_param_name = m._param_name_groups[i][0]
                         names = [backing_param_name]
                         shapes = [p._orig_size]
                         numels = [p._orig_size.numel()]

From 1be7aa035e247851bae0a15db6f383c395b8e9e3 Mon Sep 17 00:00:00 2001
From: Shikai Li <shikaili@fb.com>
Date: Mon, 8 Apr 2024 13:57:14 -0700
Subject: [PATCH 07/20] Fix no pp hanging error.

- Cleans up amax and scale update logic. Amax and scale should be
  done for both weights and parameters. So it should be done at
  forward of each microbatch.

- Consolidate `cast_params` and `all_gather` stream.
---
 .../fully_sharded_data_parallel.py            | 169 +++++++++++-------
 fairscale/nn/misc/flatten_params_wrapper.py   |   3 +-
 2 files changed, 103 insertions(+), 69 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index f3fa2acb8..2f370bee6 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -589,10 +589,9 @@ def should_flatten(name: str) -> bool:
         self.dont_wait_current_stream_for_post_all_gather = False
         self._all_gather_free_event_queue = _FreeEventQueue() if limit_all_gather_events else None
         self._reduce_scatter_free_event_queue = _FreeEventQueue() if limit_reduce_scatter_events else None
-        self.is_full_params_first_rebuilt = True
 
     @property
-    def _is_fp8_compute(self) -> bool:
+    def fp8_compute(self) -> bool:
         return _is_fp8_dtype(self.compute_dtype)
 
     def _get_gradient_predivide_factor(self, world_size: int) -> float:
@@ -832,7 +831,7 @@ def _shard_parameters_(self) -> None:
 
             # If world_size is 1, then we all-reduce grads instead of sharding.
             p._is_sharded = (self.world_size > 1) and (
-                not self._is_fp8_compute or isinstance(p, FlatParameter)
+                not self.fp8_compute or isinstance(p, FlatParameter)
             )
             p._orig_size = p.data.size()
 
@@ -1193,7 +1192,11 @@ def summon_full_params(self, recurse: bool = True, volatile: bool = False) -> Ge
             # Set the state so that we assert when trying to go into
             # forward/backward.
             self.training_state = TrainingState.SUMMON_FULL_PARAMS
-            full_tensors = self._rebuild_full_params(force_full_precision=True)
+            full_tensors = self._rebuild_full_params(
+                force_full_precision=True,
+                wait_for_all_gather=True,
+                is_first_microbatch_fwd=False,
+            )
             assert full_tensors is not None
             with contextlib.ExitStack() as stack:
                 if self.module.is_flattened:
@@ -1286,8 +1289,8 @@ def _param_dtype(self, p: Parameter) -> torch.dtype:
         Returns:
             The dtype to use for the sharded parameters.
         """
-        if self._is_fp8_compute and (not self.fp8_all_gather or
-                                     not isinstance(p, FlatParameter)):
+        if self.fp8_compute and (not self.fp8_all_gather or
+                                 not isinstance(p, FlatParameter)):
             return torch.bfloat16
         else:
             return self.compute_dtype
@@ -1435,8 +1438,6 @@ def _setup_streams(self) -> None:
             return
 
         if torch.cuda.is_available():
-            # Stream to move main FP32 params (may be on CPU) to FP32/FP16/FP8 for forward.
-            self._streams["cast_param"] = torch.cuda.Stream()
             # Stream for all-gathering parameters.
             self._streams["all_gather"] = torch.cuda.Stream()
             # Stream for overlapping grad reduction with the backward pass.
@@ -1471,10 +1472,7 @@ def _wait_for_previous_optim_step(self) -> None:
         """
         if not torch.cuda.is_available():
             return
-        if self.mixed_precision or self.move_params_to_cpu:
-            self._streams["cast_param"].wait_stream(torch.cuda.current_stream())
-        else:
-            self._streams["all_gather"].wait_stream(torch.cuda.current_stream())
+        self._streams["all_gather"].wait_stream(torch.cuda.current_stream())
 
     def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
         self._lazy_init()
@@ -1484,7 +1482,7 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
 
         # For root and mixed precision, we convert the input to FP16 (no_grad is needed for
         # the conversion).
-        is_bf16 = (self.compute_dtype == torch.bfloat16) or self._is_fp8_compute
+        is_bf16 = (self.compute_dtype == torch.bfloat16) or self.fp8_compute
         if self._is_root and self.mixed_precision:
             args, kwargs = cast_floats_to_right_precision(True, True, is_bf16, *args, **kwargs)
 
@@ -1498,20 +1496,23 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
         if self.force_input_to_fp32 and not self.mixed_precision:
             args, kwargs = cast_floats_to_right_precision(False, False, is_bf16, *args, **kwargs)
 
-        self.module.is_first_batch = not getattr(self, "is_not_first_batch", False)
-        self.is_not_first_batch = True
-
         # All-gather full parameters. This will also transfer FP32 parameters to
         # ``self.compute_dtype`` (e.g., FP16 if *mixed_precision* is ``True``).
-        self._rebuild_full_params()
-        self.is_full_params_first_rebuilt = False
+        self.module.has_unflatten_views = getattr(self.module, "has_unflatten_views", False)
 
+        is_first_microbatch_fwd=kwargs.get("is_first_microbatch", True)
+        self._rebuild_full_params(
+            wait_for_all_gather=True,
+            is_first_microbatch_fwd=is_first_microbatch_fwd
+        )
         if (
             self._fsdp_forward_ordering is not None
-            and self._my_fsdp_instance_idx is not None and self._my_fsdp_instance_idx < len(self._fsdp_forward_ordering) - 1
+            and self._my_fsdp_instance_idx is not None
+            and self._my_fsdp_instance_idx < len(self._fsdp_forward_ordering) - 1
         ):
             self._fsdp_forward_ordering[self._my_fsdp_instance_idx + 1]._rebuild_full_params(
-                wait_for_all_gather=False
+                wait_for_all_gather=False,
+                is_first_microbatch_fwd=is_first_microbatch_fwd
             )
 
         # Register backward hooks to reshard params and reduce-scatter grads.
@@ -1602,13 +1603,18 @@ def _pre_backward_hook(*unused: Any) -> None:
             # idempotent.  So in case they are called unnecessarily, they don't incur much
             # overhead.
             if self.reshard_after_forward:
-                self._rebuild_full_params()
+                self._rebuild_full_params(
+                    wait_for_all_gather=True,
+                    is_first_microbatch_fwd=False,
+                )
                 if (
-                    self.reshard_after_forward
-                    and self._fsdp_forward_ordering is not None
+                    self._fsdp_forward_ordering is not None
                     and self._my_fsdp_instance_idx is not None and self._my_fsdp_instance_idx > 0
                 ):
-                    self._fsdp_forward_ordering[self._my_fsdp_instance_idx - 1]._rebuild_full_params(wait_for_all_gather=False)
+                    self._fsdp_forward_ordering[self._my_fsdp_instance_idx - 1]._rebuild_full_params(
+                        wait_for_all_gather=False,
+                        is_first_microbatch_fwd=False,
+                    )
             else:
                 self._use_full_params()
 
@@ -1787,6 +1793,7 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
         main_grad = getattr(param, "main_grad", None)
         # Only one of `grad` or `main_grad` can exists. Whenever `main_grad is used for accumulation,
         # grad should be set as `None`.
+        assert not param.requires_grad or (grad is not None or main_grad is not None)
         assert not (grad is not None and main_grad is not None)
         # Use `grad` or `main_grad` that is not None and avoid invoking a kernel to check all zeros.
         to_reduce_grad = grad if grad is not None else main_grad
@@ -1803,7 +1810,7 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
             # Free full params.
             self._free_full_params([param])
 
-        if self.mixed_precision and (self._require_backward_grad_sync or self.reshard_after_forward) and not self.fp8_all_gather:
+        if self.mixed_precision and (self._require_backward_grad_sync or self.reshard_after_forward):
             # This is a no-op if reshard_after_forward is True, since we already
             # free the param shard when rebuilding the full params in the
             # pre_backward_hook.
@@ -1876,7 +1883,7 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                 )
             else:
                 # Unsharded parameters only happens for word_size == 1 or self.fp8_all_gather
-                assert self.world_size == 1 or (self._is_fp8_compute and self.fp8_all_gather and not isinstance(param, FlatParameter))
+                assert self.world_size == 1 or (self.fp8_compute and self.fp8_all_gather and not isinstance(param, FlatParameter))
                 if self.world_size > 1:
                     torch.distributed.all_reduce(
                         to_reduce_grad,
@@ -1896,7 +1903,6 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                 release_full_grad_event.record()
                 self._reduce_scatter_free_event_queue.enqueue(release_full_grad_event)
 
-
     def _post_reduction_hook(self, param: Parameter, reduced_grad: torch.Tensor) -> None:
         """Hook to call on each param after the reduce-scatter."""
         assert torch.cuda.current_stream() == self._streams["post_backward"]
@@ -2103,7 +2109,8 @@ def _rebuild_full_params_recursive(self):
             if isinstance(module, FullyShardedDataParallel):
                 module._lazy_init()
                 module._rebuild_full_params(
-                    wait_for_all_gather=False, not_from_recursive=False
+                    wait_for_all_gather=False,
+                    is_first_microbatch_fwd=True,
                 )
 
 
@@ -2112,8 +2119,8 @@ def _rebuild_full_params_recursive(self):
     def _rebuild_full_params(
         self,
         force_full_precision: bool = False,
-        wait_for_all_gather=True,
-        not_from_recursive=True,
+        wait_for_all_gather: bool = True,
+        is_first_microbatch_fwd: bool = False,
     ) -> Optional[List[Tuple[torch.Tensor, bool]]]:
         """
         Gather all shards of params.
@@ -2169,6 +2176,9 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
 
             self.has_full_params = False
 
+        if self.fp8_compute and self.fp8_all_gather:
+            self._update_amax_and_scale_fwd(is_first_microbatch_fwd=is_first_microbatch_fwd)
+
         if self._has_shared_params:
             # self.has_full_params flag can be out of sync if a shared param is
             # sharded by another FSDP instance. An example is that in eval case
@@ -2203,7 +2213,7 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
             if (
                 self.mixed_precision or self.move_params_to_cpu
             ) and not force_full_precision:
-                self._cast_params_for_all_gather(not_from_recursive=not_from_recursive)
+                self._cast_params_for_all_gather()
 
             if self.move_params_to_cpu:
                 if force_full_precision:
@@ -2211,7 +2221,7 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
                     # use pinned memory. Otherwise move p.data to the compute
                     # device.
                     if self.params[0].dtype == self.compute_dtype:
-                        self._cast_params_for_all_gather(not_from_recursive=not_from_recursive)
+                        self._cast_params_for_all_gather()
                     else:
                         for p in self.params:
                             p.data = p.data.to(self.compute_device)
@@ -2325,8 +2335,8 @@ def _free_full_params(self, params: Optional[List[Parameter]] = None) -> None:
         """Free up storage for full parameters."""
         if params is None:
             params = self.params
-        self.is_not_first_batch = False
-        self.is_full_params_first_rebuilt = True
+
+        self.module.has_unflatten_views = False
         self.has_full_params = False
 
         current_stream = torch.cuda.current_stream()
@@ -2481,25 +2491,21 @@ def _use_fp32_param_shard(self, params: Optional[List[Parameter]] = None) -> Non
             p.data = p._fp32_shard
 
     @torch.no_grad()
-    def _cast_params_for_all_gather(
-        self, params: Optional[List[Parameter]] = None,
-        not_from_recursive: bool = False,
-    ) -> None:
-        """Cast FP32 param shard to FP16 for a list of params."""
+    def _update_amax_and_scale_fwd(
+        self,
+        params: Optional[List[Parameter]] = None,
+        is_first_microbatch_fwd: bool = False,
+    ):
+        """Update Amax and scales associated with FP8 parameters."""
         if params is None:
             params = self.params
 
-        with torch.cuda.stream(self._streams["cast_param"]):
+        with torch.cuda.stream(self._streams["all_gather"]):
             for p in params:
-                assert p._fp16_shard is not None
-                alloc_storage_(p._fp16_shard, size=p._fp32_shard.size())
-
-                if self._is_fp8_compute and _is_fp8_dtype(p._fp16_shard.dtype):
+                if self.fp8_compute and _is_fp8_dtype(p._fp16_shard.dtype):
                     assert isinstance(p, FlatParameter)
                     assert len(p._param_infos) == len(p._param_numels)
 
-                    numel_per_shard = p.numel()
-                    offset = -numel_per_shard * self.rank
                     for i in range(len(p._param_infos)):
                         _, m, n = p._param_infos[i]
                         assert _is_te_module_with_weights(m)
@@ -2509,27 +2515,55 @@ def _cast_params_for_all_gather(
                                 num_gemms=2 if isinstance(m, te.LayerNormMLP) else 1
                             )
 
-                        if self.is_full_params_first_rebuilt:
-                            if m.fp8_meta.get("update_amax_and_scale_fwd", False):
-                                if m.fp8_meta["recipe"].reduce_amax:
-                                    FP8GlobalStateManager.copy_amax_from_global_buffer(
-                                        m.fp8_meta, forward=True
-                                    )
-                                    amax_and_scale_update(
-                                        m.fp8_meta,
-                                        True,
-                                        update_weight_scale_inv=True,
-                                    )
-                                    if not_from_recursive:
-                                        FP8GlobalStateManager.set_amax_buffer_key_deletion(
-                                            m.fp8_meta, forward=True
-                                        )
-                                else:
-                                    amax_and_scale_update(
-                                        m.fp8_meta,
-                                        True,
-                                        update_weight_scale_inv=True,
-                                    )
+                        if m.fp8_meta.get("update_amax_and_scale_fwd", False):
+                            if m.fp8_meta["recipe"].reduce_amax:
+                                logging.warning(f"Reduce amax! {is_first_microbatch_fwd}")
+                                FP8GlobalStateManager.copy_amax_from_global_buffer(
+                                    m.fp8_meta, forward=True
+                                )
+                                amax_and_scale_update(
+                                    m.fp8_meta,
+                                    True,
+                                    update_weight_scale_inv=is_first_microbatch_fwd,
+                                )
+                                FP8GlobalStateManager.set_amax_buffer_key_deletion(
+                                    m.fp8_meta, forward=True
+                                )
+                            else:
+                                logging.warning("Not reduce amax! {is_first_microbatch_fwd}")
+                                amax_and_scale_update(
+                                    m.fp8_meta,
+                                    True,
+                                    update_weight_scale_inv=is_first_microbatch_fwd,
+                                )
+                            m.fp8_meta["update_amax_and_scale_fwd"] = False
+
+
+
+    @torch.no_grad()
+    def _cast_params_for_all_gather(
+        self,
+        params: Optional[List[Parameter]] = None,
+    ) -> None:
+        """Cast FP32 params shard to FP16/BF16/FP8 for a list of params."""
+        if params is None:
+            params = self.params
+
+        with torch.cuda.stream(self._streams["all_gather"]):
+            for p in params:
+                assert p._fp16_shard is not None
+                alloc_storage_(p._fp16_shard, size=p._fp32_shard.size())
+
+                if self.fp8_compute and _is_fp8_dtype(p._fp16_shard.dtype):
+                    assert isinstance(p, FlatParameter), "FP8 parameters should be all flatten"
+                    assert len(p._param_infos) == len(p._param_numels)
+
+                    numel_per_shard = p.numel()
+                    offset = -numel_per_shard * self.rank
+                    for i in range(len(p._param_infos)):
+                        _, m, n = p._param_infos[i]
+                        assert _is_te_module_with_weights(m), "Modules with FP8 parameters shoule be TE modules"
+                        assert m.fp8_initialized, "Modules with FP8 parameters should be initialized with scales"
 
                         numel = p._param_numels[i]
                         if offset + numel <= 0 or offset >= numel_per_shard:
@@ -2566,7 +2600,6 @@ def _cast_params_for_all_gather(
                     )
                     p.data = p._fp16_shard
 
-        self._streams["all_gather"].wait_stream(self._streams["cast_param"])
 
     @torch.no_grad()
     def _free_fp16_param_shard(self, params: Optional[List[Parameter]] = None) -> None:
diff --git a/fairscale/nn/misc/flatten_params_wrapper.py b/fairscale/nn/misc/flatten_params_wrapper.py
index ae2f3d792..da947dc31 100644
--- a/fairscale/nn/misc/flatten_params_wrapper.py
+++ b/fairscale/nn/misc/flatten_params_wrapper.py
@@ -486,8 +486,9 @@ def load_state_dict(
             return super().load_state_dict(state_dict, strict)
 
     def forward(self, *inputs: Any, **kwinputs: Any) -> Any:
-        if getattr(self, "is_first_batch", False):
+        if not getattr(self, "has_unflatten_views", False):
             self._unflatten_params_as_views()
+            self.has_unflatten_views = True
         return self.module(*inputs, **kwinputs)
 
     def get_param_views(self, external_data_list: Optional[List[Optional[Tensor]]] = None) -> Iterator[Tensor]:

From 57eb557a692b795df377983e6dda36eeb35cf1ae Mon Sep 17 00:00:00 2001
From: Shikai Li <shikaili@fb.com>
Date: Wed, 10 Apr 2024 13:16:41 -0700
Subject: [PATCH 08/20] Clean up shard offset calculation logic.

---
 .../fully_sharded_data_parallel.py            | 30 ++++++++++++-------
 1 file changed, 19 insertions(+), 11 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index 2f370bee6..99e07899b 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -440,7 +440,7 @@ def __init__(
         self.flatten_parameters = flatten_parameters
         self.move_params_to_cpu = move_params_to_cpu or cpu_offload
         self.compute_dtype = compute_dtype or (torch.float16 if mixed_precision else torch.float32)
-        self.buffer_dtype = buffer_dtype or self.compute_dtype
+        self.buffer_dtype = buffer_dtype or (torch.bfloat16 if _is_fp8_dtype(self.compute_dtype) else self.compute_dtype)
         self.move_grads_to_cpu = self.move_params_to_cpu if move_grads_to_cpu is None else move_grads_to_cpu
         self.bucket_cap_mb = bucket_cap_mb
         self.compute_device = compute_device or _get_default_cuda_device(module)
@@ -2517,7 +2517,6 @@ def _update_amax_and_scale_fwd(
 
                         if m.fp8_meta.get("update_amax_and_scale_fwd", False):
                             if m.fp8_meta["recipe"].reduce_amax:
-                                logging.warning(f"Reduce amax! {is_first_microbatch_fwd}")
                                 FP8GlobalStateManager.copy_amax_from_global_buffer(
                                     m.fp8_meta, forward=True
                                 )
@@ -2530,7 +2529,6 @@ def _update_amax_and_scale_fwd(
                                     m.fp8_meta, forward=True
                                 )
                             else:
-                                logging.warning("Not reduce amax! {is_first_microbatch_fwd}")
                                 amax_and_scale_update(
                                     m.fp8_meta,
                                     True,
@@ -2555,28 +2553,38 @@ def _cast_params_for_all_gather(
                 alloc_storage_(p._fp16_shard, size=p._fp32_shard.size())
 
                 if self.fp8_compute and _is_fp8_dtype(p._fp16_shard.dtype):
+                    assert p._is_sharded
                     assert isinstance(p, FlatParameter), "FP8 parameters should be all flatten"
                     assert len(p._param_infos) == len(p._param_numels)
 
                     numel_per_shard = p.numel()
-                    offset = -numel_per_shard * self.rank
+
+                    flat_index = 0
+                    flat_begin = numel_per_shard * self.rank
+                    flat_end = flat_begin + numel_per_shard
+
                     for i in range(len(p._param_infos)):
                         _, m, n = p._param_infos[i]
+
                         assert _is_te_module_with_weights(m), "Modules with FP8 parameters shoule be TE modules"
                         assert m.fp8_initialized, "Modules with FP8 parameters should be initialized with scales"
 
                         numel = p._param_numels[i]
-                        if offset + numel <= 0 or offset >= numel_per_shard:
-                            offset += numel
+
+                        if flat_index >= flat_end:
+                            break
+                        shard_begin = max(flat_index - flat_begin, 0)
+
+                        flat_index += numel
+                        if flat_index <= flat_begin:
                             continue
+                        shard_end = min(flat_index - flat_begin, numel_per_shard)
 
                         fp8_dtype_forward = te.fp8.get_fp8_te_dtype(
                             m.fp8_meta["recipe"], fprop_tensor=True
                         )
-                        begin = max(offset, 0)
-                        end = min(offset + numel, numel_per_shard)
                         cast_to_fp8(
-                            p._fp32_shard[begin:end].bfloat16(),
+                            p._fp32_shard[shard_begin:shard_end].bfloat16().contiguous(),
                             m.fp8_meta["scaling_fwd"],
                             (
                                 FP8FwdTensors.GEMM2_WEIGHT
@@ -2584,9 +2592,9 @@ def _cast_params_for_all_gather(
                                 else FP8FwdTensors.GEMM1_WEIGHT
                             ),
                             fp8_dtype_forward,
-                            out=p._fp16_shard[begin:end],
+                            out=p._fp16_shard[shard_begin:shard_end],
                         )
-                        offset += numel
+                    # Doesn't need to set padding elements.
                     p.data = p._fp16_shard.view(
                             torch.float8_e4m3fn
                             if fp8_dtype_forward == DType.kFloat8E4M3

From e9e8f8ec854fe00e417ec7124ad59bc0ff5f87b9 Mon Sep 17 00:00:00 2001
From: Shikai Li <shikaili@fb.com>
Date: Mon, 15 Apr 2024 23:20:47 -0700
Subject: [PATCH 09/20] Unify compute dtype setting.

---
 .../fully_sharded_data_parallel.py            | 82 ++++++++++---------
 1 file changed, 44 insertions(+), 38 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index 99e07899b..3de055a73 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -439,8 +439,8 @@ def __init__(
         self.fp8_all_gather = fp8_all_gather
         self.flatten_parameters = flatten_parameters
         self.move_params_to_cpu = move_params_to_cpu or cpu_offload
-        self.compute_dtype = compute_dtype or (torch.float16 if mixed_precision else torch.float32)
-        self.buffer_dtype = buffer_dtype or (torch.bfloat16 if _is_fp8_dtype(self.compute_dtype) else self.compute_dtype)
+        self.compute_dtype = compute_dtype or (torch.bfloat16 if mixed_precision else torch.float32)
+        self.buffer_dtype = buffer_dtype or self.compute_dtype
         self.move_grads_to_cpu = self.move_params_to_cpu if move_grads_to_cpu is None else move_grads_to_cpu
         self.bucket_cap_mb = bucket_cap_mb
         self.compute_device = compute_device or _get_default_cuda_device(module)
@@ -485,6 +485,11 @@ def __init__(
                 param_names.append(param_name)
                 params.append(param)
 
+        for m in module.modules():
+            for param in m.parameters():
+                if not getattr(param, "_is_te_param", False):
+                    param._is_te_param = _is_te_module_with_weights(m)
+
         self._has_params = len(params) > 0
         self._has_shared_params = False
 
@@ -501,27 +506,34 @@ def __init__(
 
         # For now, it is either all flatten or none flatten. This will be extended to
         # multiple flatten groups in my next PR.
-        def should_flatten(name: str) -> bool:
+        no_te_params = not any(p._is_te_param for p in params)
+        def should_flatten(name: str, param: Parameter) -> bool:
+            if not self.flatten_parameters:
+                return False
+            # If no TE weights or no FP8 AllGather, then flatten them all. Shard as compute dtype.
+            if no_te_params or not self.fp8_all_gather:
+                return True
             # `*_norm_weights` are numerics-sensitive and cannot be quantized to fp8.
-            return self.flatten_parameters and (not self.fp8_all_gather or "norm_weight" not in name)
+            return  param._is_te_param and "norm_weight" not in name
+
+        to_be_flatten_param_names = []
+        to_be_flatten_params = []
+        non_flatten_param_names = []
+        non_flatten_params = []
+        for name, param in zip(param_names, params):
+            if should_flatten(name, param):
+                to_be_flatten_param_names.append(name)
+                to_be_flatten_params.append(param)
+            else:
+                non_flatten_param_names.append(name)
+                non_flatten_params.append(param)
 
-        to_be_flatten_params: List[List[Parameter]] = [
-            param
-            for param, name in zip(params, param_names)
-            if should_flatten(name)
-        ]
         if to_be_flatten_params:
             to_be_flatten_params = [to_be_flatten_params]
-        non_flatten_params: List[List[Parameter]] = [
-            param
-            for param, name in zip(params, param_names)
-            if not should_flatten(name)
-        ]
-        param_name_groups: List[List[str]] = [
-            [n for n in param_names if should_flatten(n)]
-        ] + [
-            [n] for n in param_names if not should_flatten(n)
-        ]
+
+        param_name_groups: List[List[str]] = [to_be_flatten_param_names] + [[n] for n in non_flatten_param_names]
+
+        logging.info(f"param_names: {param_name_groups}")
         del param_names
 
         self._fsdp_wrapped_module: nn.Module = FlattenParamsWrapper(
@@ -529,6 +541,9 @@ def should_flatten(name: str) -> bool:
         )
         del module  # free original module in case it helps garbage collection
 
+        for param in self._fsdp_wrapped_module.flat_params:
+            param._is_fp8_param = not no_te_params and self.fp8_all_gather
+
         # Now, in this FSDP wrapper class, we keep a list of to-be-flatten and not-to-be-flatten
         # params for doing sharding, gradient hooks, etc. Note, the ordering of the
         # list matters: flatten params are always in the front.
@@ -590,10 +605,6 @@ def should_flatten(name: str) -> bool:
         self._all_gather_free_event_queue = _FreeEventQueue() if limit_all_gather_events else None
         self._reduce_scatter_free_event_queue = _FreeEventQueue() if limit_reduce_scatter_events else None
 
-    @property
-    def fp8_compute(self) -> bool:
-        return _is_fp8_dtype(self.compute_dtype)
-
     def _get_gradient_predivide_factor(self, world_size: int) -> float:
         factor: int = 1
         while world_size % factor == 0 and world_size / factor > factor:
@@ -830,9 +841,7 @@ def _shard_parameters_(self) -> None:
                 assert p.dtype == torch.float32
 
             # If world_size is 1, then we all-reduce grads instead of sharding.
-            p._is_sharded = (self.world_size > 1) and (
-                not self.fp8_compute or isinstance(p, FlatParameter)
-            )
+            p._is_sharded = (self.world_size > 1) and isinstance(p, FlatParameter)
             p._orig_size = p.data.size()
 
             if not p._is_sharded:
@@ -1289,11 +1298,9 @@ def _param_dtype(self, p: Parameter) -> torch.dtype:
         Returns:
             The dtype to use for the sharded parameters.
         """
-        if self.fp8_compute and (not self.fp8_all_gather or
-                                 not isinstance(p, FlatParameter)):
-            return torch.bfloat16
-        else:
-            return self.compute_dtype
+        if getattr(p, "_is_fp8_param", False):
+            return torch.float8_e4m3fn
+        return self.compute_dtype
 
     @torch.no_grad()
     def _init_param_attributes(self, p: Parameter) -> None:
@@ -1482,7 +1489,7 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
 
         # For root and mixed precision, we convert the input to FP16 (no_grad is needed for
         # the conversion).
-        is_bf16 = (self.compute_dtype == torch.bfloat16) or self.fp8_compute
+        is_bf16 = self.compute_dtype == torch.bfloat16
         if self._is_root and self.mixed_precision:
             args, kwargs = cast_floats_to_right_precision(True, True, is_bf16, *args, **kwargs)
 
@@ -1514,7 +1521,6 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
                 wait_for_all_gather=False,
                 is_first_microbatch_fwd=is_first_microbatch_fwd
             )
-
         # Register backward hooks to reshard params and reduce-scatter grads.
         # These need to be re-registered every forward pass.
         self._register_post_backward_hooks()
@@ -1883,7 +1889,7 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                 )
             else:
                 # Unsharded parameters only happens for word_size == 1 or self.fp8_all_gather
-                assert self.world_size == 1 or (self.fp8_compute and self.fp8_all_gather and not isinstance(param, FlatParameter))
+                assert self.world_size == 1 or not isinstance(param, FlatParameter)
                 if self.world_size > 1:
                     torch.distributed.all_reduce(
                         to_reduce_grad,
@@ -2176,7 +2182,7 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
 
             self.has_full_params = False
 
-        if self.fp8_compute and self.fp8_all_gather:
+        if self.fp8_all_gather:
             self._update_amax_and_scale_fwd(is_first_microbatch_fwd=is_first_microbatch_fwd)
 
         if self._has_shared_params:
@@ -2502,7 +2508,7 @@ def _update_amax_and_scale_fwd(
 
         with torch.cuda.stream(self._streams["all_gather"]):
             for p in params:
-                if self.fp8_compute and _is_fp8_dtype(p._fp16_shard.dtype):
+                if _is_fp8_dtype(p._fp16_shard.dtype):
                     assert isinstance(p, FlatParameter)
                     assert len(p._param_infos) == len(p._param_numels)
 
@@ -2552,7 +2558,7 @@ def _cast_params_for_all_gather(
                 assert p._fp16_shard is not None
                 alloc_storage_(p._fp16_shard, size=p._fp32_shard.size())
 
-                if self.fp8_compute and _is_fp8_dtype(p._fp16_shard.dtype):
+                if _is_fp8_dtype(p._fp16_shard.dtype):
                     assert p._is_sharded
                     assert isinstance(p, FlatParameter), "FP8 parameters should be all flatten"
                     assert len(p._param_infos) == len(p._param_numels)
@@ -2566,7 +2572,7 @@ def _cast_params_for_all_gather(
                     for i in range(len(p._param_infos)):
                         _, m, n = p._param_infos[i]
 
-                        assert _is_te_module_with_weights(m), "Modules with FP8 parameters shoule be TE modules"
+                        assert _is_te_module_with_weights(m), f"Modules {m} with FP8 parameters shoule be TE modules"
                         assert m.fp8_initialized, "Modules with FP8 parameters should be initialized with scales"
 
                         numel = p._param_numels[i]

From 21f8e051e49813d7cd09408f63e950f51417f528 Mon Sep 17 00:00:00 2001
From: Shikai Li <shikaili@fb.com>
Date: Wed, 17 Apr 2024 16:45:15 -0700
Subject: [PATCH 10/20] Have FP16 and FP8 sharded in the same way.

---
 fairscale/nn/data_parallel/fully_sharded_data_parallel.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index 3de055a73..22e6261b9 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -511,7 +511,7 @@ def should_flatten(name: str, param: Parameter) -> bool:
             if not self.flatten_parameters:
                 return False
             # If no TE weights or no FP8 AllGather, then flatten them all. Shard as compute dtype.
-            if no_te_params or not self.fp8_all_gather:
+            if no_te_params: #or not self.fp8_all_gather:
                 return True
             # `*_norm_weights` are numerics-sensitive and cannot be quantized to fp8.
             return  param._is_te_param and "norm_weight" not in name

From 8ec7c1df8ecb1e963fe40960bbe2ebe9cd8f3d07 Mon Sep 17 00:00:00 2001
From: ngoyal2707 <ngoyal2707@users.noreply.github.com>
Date: Wed, 24 Jan 2024 15:49:34 -0500
Subject: [PATCH 11/20] added option for no PG validation for faster init
 (#1161)

Co-authored-by: Naman Goyal <naman@fb.com>
---
 fairscale/nn/data_parallel/fully_sharded_data_parallel.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index bfa02ae77..662ece626 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -371,6 +371,7 @@ def __init__(
         limit_all_gather_events: bool = False,
         limit_reduce_scatter_events: bool = False,
         cast_input: bool = True,
+        should_validate_process_group: bool = True,
     ):
         try:
             import torch._C
@@ -455,7 +456,7 @@ def __init__(
             raise ValueError(f"offload type: '{offload_config.offload_type}' requires flatten_parameters=True")
 
         # skip validation if the process group was created above
-        if process_group:
+        if process_group and should_validate_process_group:
             validate_process_group(self.compute_device, self.process_group)
 
         # enable pytorch sync_bn just in case model contains sync_bn layers.

From f27ab176c3ac9fd21c545f07c6843eab13c9a44b Mon Sep 17 00:00:00 2001
From: Shikai Li <shikaili@fb.com>
Date: Tue, 12 Mar 2024 19:52:59 -0700
Subject: [PATCH 12/20] Mirros Jiecao's change.

---
 .../fully_sharded_data_parallel.py            | 223 ++++++++++++++++--
 fairscale/nn/misc/__init__.py                 |   2 +-
 fairscale/nn/misc/flatten_params_wrapper.py   |   3 +-
 3 files changed, 201 insertions(+), 27 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index 662ece626..ca744cafe 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -44,8 +44,11 @@
 import torch.nn.functional as F
 from torch.nn.parameter import Parameter
 from torch.utils.hooks import RemovableHandle
+import transformer_engine.pytorch as te
+from transformer_engine.pytorch.cpp_extensions import cast_to_fp8, DType, FP8FwdTensors
+from transformer_engine.pytorch.fp8 import amax_and_scale_update, FP8GlobalStateManager
 
-from fairscale.nn.misc import FlattenParamsWrapper
+from fairscale.nn.misc import FlatParameter, FlattenParamsWrapper
 from fairscale.nn.wrap import auto_wrap, config_auto_wrap_policy, enable_wrap
 from fairscale.utils.containers import apply_to_tensors
 from fairscale.utils.parallel import (
@@ -152,6 +155,14 @@ class OffloadConfig:
     dir: Optional[str] = None
 
 
+def _is_fp8_dtype(dtype: torch.dtype) -> bool:
+    return dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
+
+
+def _is_te_module_with_weights(m: nn.Module) -> bool:
+    return isinstance(m, (te.Linear, te.LayerNormLinear, te.LayerNormMLP))
+
+
 class FullyShardedDataParallel(nn.Module):
     """
     A wrapper for sharding Module parameters across data parallel workers. This
@@ -372,6 +383,7 @@ def __init__(
         limit_reduce_scatter_events: bool = False,
         cast_input: bool = True,
         should_validate_process_group: bool = True,
+        fp8_allgather: bool = False,
     ):
         try:
             import torch._C
@@ -438,6 +450,7 @@ def __init__(
         self.force_input_to_fp32 = force_input_to_fp32
         self.verbose = verbose
         self.state_dict_on_rank_0_only = state_dict_on_rank_0_only
+        self.fp8_allgather = fp8_allgather
         # Experimental feature for now. Use at your own risk.
         self.ssd_offload = True if offload_config and offload_config.offload_type == "ssd_offload" else False
 
@@ -492,9 +505,22 @@ def __init__(
         non_flatten_params = params
         param_name_groups = [[n] for n in param_names]
         if self.flatten_parameters:
-            to_be_flatten_params = [params]
-            non_flatten_params = []
-            param_name_groups = [param_names]
+            to_be_flatten_params = [
+                [
+                    params[i]
+                    for i in range(len(params))
+                    if "norm_weight" not in param_names[i]
+                ]
+            ]
+            non_flatten_params = [
+                params[i]
+                for i in range(len(params))
+                if "norm_weight" in param_names[i]
+            ]
+            param_name_groups = [
+                [n for n in param_names if "norm_weight" not in n],
+                [n for n in param_names if "norm_weight" in n],
+            ]
         del param_names
 
         self._fsdp_wrapped_module: nn.Module = FlattenParamsWrapper(
@@ -563,6 +589,10 @@ def __init__(
         self._all_gather_free_event_queue = _FreeEventQueue() if limit_all_gather_events else None
         self._reduce_scatter_free_event_queue = _FreeEventQueue() if limit_reduce_scatter_events else None
 
+    @property
+    def _is_fp8_compute(self) -> bool:
+        return _is_fp8_dtype(self.compute_dtype)
+
     def _get_gradient_predivide_factor(self, world_size: int) -> float:
         factor: int = 1
         while world_size % factor == 0 and world_size / factor > factor:
@@ -795,7 +825,9 @@ def _shard_parameters_(self) -> None:
                 assert p.dtype == torch.float32
 
             # If world_size is 1, then we all-reduce grads instead of sharding.
-            p._is_sharded = self.world_size > 1
+            p._is_sharded = (self.world_size > 1) and (
+                not self._is_fp8_compute or isinstance(p, FlatParameter)
+            )
             p._orig_size = p.data.size()
 
             if not p._is_sharded:
@@ -1179,7 +1211,7 @@ def summon_full_params(self, recurse: bool = True, volatile: bool = False) -> Ge
                         non_shared_params
                     ), f"{len(full_tensors)} vs. {len(non_shared_params)}"
                     for p, (full_tensor, safe_to_free) in zip(non_shared_params, full_tensors):
-                        if not volatile:
+                        if not volatile and p._is_sharded:
                             # Copy any changes made to the full params back into
                             # the corresponding local shards.
                             local_shard, _ = self._get_shard(full_tensor)
@@ -1241,6 +1273,20 @@ def _lazy_init(self) -> None:
             # ``optim.step()`` is done before we all-gather parameters.
             self._wait_for_previous_optim_step()
 
+    def _shard_dtype(self, p: Parameter) -> torch.dtype:
+        """
+        Return the dtype to use for the sharded parameters.
+
+        Returns:
+            The dtype to use for the sharded parameters.
+        """
+        if self._is_fp8_compute and not isinstance(p, FlatParameter):
+            # Assume non flattened are precision critical like norm
+            assert not p._is_sharded
+            return torch.bfloat16
+        else:
+            return self.compute_dtype
+
     @torch.no_grad()
     def _init_param_attributes(self, p: Parameter) -> None:
         """
@@ -1299,7 +1345,9 @@ def _init_param_attributes(self, p: Parameter) -> None:
             # storage to size 0 at init (here) and re-materialize (by copying
             # from _fp32_shard) as needed. If offloading params to CPU, the
             # dtype of the fp16 shard will depend on the *`compute_dtype`*.
-            p._fp16_shard = torch.zeros_like(p._fp32_shard, device=self.compute_device, dtype=self.compute_dtype)
+            p._fp16_shard = torch.zeros_like(
+                p._fp32_shard, device=self.compute_device, dtype=self._shard_dtype(p)
+            )
             free_storage_(p._fp16_shard)
 
         if self.mixed_precision:
@@ -1318,7 +1366,9 @@ def _init_param_attributes(self, p: Parameter) -> None:
         # relevant computation.
         if p._is_sharded:
             p._full_param_padded = torch.zeros(
-                p.data.numel() * self.world_size, device=self.compute_device, dtype=self.compute_dtype
+                p.data.numel() * self.world_size,
+                device=self.compute_device,
+                dtype=self._shard_dtype(p),
             )
             free_storage_(p._full_param_padded)
 
@@ -1433,7 +1483,7 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
 
         # For root and mixed precision, we convert the input to FP16 (no_grad is needed for
         # the conversion).
-        is_bf16 = self.compute_dtype == torch.bfloat16
+        is_bf16 = (self.compute_dtype == torch.bfloat16) or self._is_fp8_compute
         if self._is_root and self.mixed_precision and self.cast_input:
             args, kwargs = cast_floats_to_right_precision(True, True, is_bf16, *args, **kwargs)
 
@@ -1447,9 +1497,13 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
         if self.force_input_to_fp32 and not self.mixed_precision:
             args, kwargs = cast_floats_to_right_precision(False, False, is_bf16, *args, **kwargs)
 
+        self.module.is_first_batch = not getattr(self, "is_not_first_batch", False)
+        self.is_not_first_batch = True
+
         # All-gather full parameters. This will also transfer FP32 parameters to
         # ``self.compute_dtype`` (e.g., FP16 if *mixed_precision* is ``True``).
         self._rebuild_full_params()
+        self.is_full_params_first_rebuilt = False
 
         if (
             self._fsdp_forward_ordering is not None
@@ -1714,8 +1768,7 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
         # then subsequent hook callbacks will see POST state.
         self.assert_state([TrainingState.BACKWARD_PRE, TrainingState.BACKWARD_POST])
         self.training_state = TrainingState.BACKWARD_POST
-        if param.grad is None:
-            return
+
 
         if hasattr(param, "_linked_param"):
             # This links to a shared param. We should finalize the linked param here.
@@ -1727,7 +1780,9 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
             if hasattr(param._linked_param, "_is_shared") and param._linked_param._is_shared:
                 param = param._linked_param
 
-        assert param.grad is not None, param.shape
+        if param.grad is None:
+            return
+
         if param.grad.requires_grad:
             raise RuntimeError("FSDP only works with gradients that don't require gradients")
 
@@ -1735,7 +1790,7 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
             # Free full params.
             self._free_full_params([param])
 
-        if self.mixed_precision and (self._require_backward_grad_sync or self.reshard_after_forward):
+        if self.mixed_precision and (self._require_backward_grad_sync or self.reshard_after_forward) and not self.fp8_allgather:
             # This is a no-op if reshard_after_forward is True, since we already
             # free the param shard when rebuilding the full params in the
             # pre_backward_hook.
@@ -1749,12 +1804,12 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                 param.unsharded_main_grad = param.grad.to(torch.float32)
             else:
                 param.unsharded_main_grad.add_(param.grad.data)
-
             param.grad = None
 
         if not self._require_backward_grad_sync:
             return
 
+
         # Wait for all work in the current stream to finish, then start the
         # reductions in post_backward stream.
         self._streams["post_backward"].wait_stream(torch.cuda.current_stream())
@@ -1801,6 +1856,19 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                 self._reducer.reduce_scatter_async(
                     grad, group=self.process_group_reduce_scatter, callback_fn=callback_fn
                 )
+            elif self.fp8_allgather:
+                if param.grad is not None:
+                    if self.fp32_reduce_scatter:
+                        param.unsharded_main_grad = param.grad.to(torch.float32)
+                    else:
+                        param.unsharded_main_grad = param.grad
+                    param.grad = None
+
+                    torch.distributed.all_reduce(
+                        param.unsharded_main_grad,
+                        group=self.process_group_reduce_scatter,
+                    )
+                    self._post_reduction_hook(param, param.unsharded_main_grad)
             else:
                 # Currently the only way for _is_sharded to be False is if
                 # world_size == 1. This could be relaxed in the future, in which
@@ -1849,8 +1917,17 @@ def _post_reduction_hook(self, param: Parameter, reduced_grad: torch.Tensor) ->
                 ), f"{param._saved_grad_shard.shape} vs {reduced_grad.shape}"
                 param._saved_grad_shard.data += reduced_grad.data
             reduced_grad = param._saved_grad_shard.data
-        elif (param.grad is None) and self.fp32_reduce_scatter:
-            param.main_grad = reduced_grad.data
+
+            if getattr(param, "unsharded_main_grad", None) is not None:
+                free_storage_(param.unsharded_main_grad.data)
+                param.unsharded_main_grad = None
+
+        elif param.grad is None:
+        # TODO(shikaili): Still having missing grad/main_grad issue.
+            if self.fp32_reduce_scatter:
+                param.main_grad = reduced_grad.data
+            else:
+                param.grad = reduced_grad.data
 
         # Optionally move gradients to CPU, typically used if one is running the optimizer on the CPU. Once the full
         # backwards pass completes, we will set `.grad` to the CPU copy.
@@ -1962,6 +2039,7 @@ def _finalize_parameters(fsdp_module: FullyShardedDataParallel) -> None:
                     # again after post-backward
                     if p.shape != p._saved_grad_shard.shape:
                         self._use_fp32_param_shard([p])
+                    assert getattr(p, "unsharded_main_grad", None) is None
                     if p._saved_grad_shard.dtype != p.dtype:
                         p.main_grad = p._saved_grad_shard
                     else:
@@ -2028,12 +2106,19 @@ def _rebuild_full_params_recursive(self):
         for module in self.modules():
             if isinstance(module, FullyShardedDataParallel):
                 module._lazy_init()
-                module._rebuild_full_params(wait_for_all_gather=False)
+                module._rebuild_full_params(
+                    wait_for_all_gather=False, not_from_recursive=False
+                )
 
 
 
     @torch.no_grad()
-    def _rebuild_full_params(self, force_full_precision: bool = False, wait_for_all_gather = True) -> Optional[List[Tuple[torch.Tensor, bool]]]:
+    def _rebuild_full_params(
+        self,
+        force_full_precision: bool = False,
+        wait_for_all_gather=True,
+        not_from_recursive=True,
+    ) -> Optional[List[Tuple[torch.Tensor, bool]]]:
         """
         Gather all shards of params.
 
@@ -2053,6 +2138,41 @@ def _rebuild_full_params(self, force_full_precision: bool = False, wait_for_all_
             caller to free the full-sized param. This will be ``None`` if
             ``force_full_precision=False`` and the full params are already gathered.
         """
+        if self._is_fp8_compute:
+            # Need to use fp32_to_fp16 stream since _cast_fp32_param_shards_to_fp16 depends on this block.
+            with torch.no_grad(), torch.cuda.stream(self._streams["fp32_to_fp16"]):
+                for p in self.params:
+                    if not isinstance(p, FlatParameter):
+                        continue
+                    d = {info[0]: info[1] for info in p._param_infos}
+                    for n, m in d.items():
+                        # Previous iteration was grad_enabled
+                        if not m.fp8_initialized:
+                            m.fp8_init(
+                                num_gemms=2 if isinstance(m, te.LayerNormMLP) else 1
+                            )
+                        if m.fp8_meta.get("update_amax_and_scale_fwd", False):
+                            if m.fp8_meta["recipe"].reduce_amax:
+                                FP8GlobalStateManager.copy_amax_from_global_buffer(
+                                    m.fp8_meta, forward=True
+                                )
+                            amax_and_scale_update(
+                                m.fp8_meta,
+                                True,
+                                update_weight_scale_inv=self.is_full_params_first_rebuilt,
+                            )
+                            if not_from_recursive:
+                                FP8GlobalStateManager.set_amax_buffer_key_deletion(
+                                    m.fp8_meta, forward=True
+                                )
+                        else:
+                            amax_and_scale_update(
+                                m.fp8_meta,
+                                True,
+                                update_weight_scale_inv=self.is_full_params_first_rebuilt,
+                            )
+                torch.cuda.current_stream().wait_stream(self._streams["fp32_to_fp16"])
+
         output_tensors: List[Tuple[torch.Tensor, bool]] = []
 
         def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
@@ -2063,6 +2183,7 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
                 custom_output_tensor (torch.Tensor, Optional): if not None, this
                 tensor contains the data we just gathered.
             """
+            p_fp16_shard_size = -1
             if custom_output_tensor is not None:
                 assert p._is_sharded
                 p.data = custom_output_tensor
@@ -2071,6 +2192,7 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
                 if (self.mixed_precision or self.move_params_to_cpu) and not force_full_precision:
                     assert p._fp16_shard is not None
                     p.data = p._fp16_shard
+                    p_fp16_shard_size = p._fp16_shard.storage().size()
                     output_tensors.append((p.data, True))
                 else:
                     # Here p.data == p._fp32_shard, so it's not safe to free.
@@ -2212,6 +2334,14 @@ def _prep_grads_for_backward(self) -> None:
         right shape, device, accumulated values, etc.
         """
         for p in self.params:
+            if isinstance(p, FlatParameter) and all(
+                _is_te_module_with_weights(info[1]) for info in p._param_infos
+            ):
+                if getattr(p, "main_grad", None) is None:
+                    p.main_grad = torch.empty_like(p, dtype=torch.float)
+                main_grad_views = p.get_param_views(p.main_grad)
+                for (_, m, n), main_grad in zip(p._param_infos, main_grad_views):
+                    getattr(m, n).main_grad = main_grad
             if p.grad is not None:
                 if p.grad.device != p.data.device:
                     p.grad = None
@@ -2233,6 +2363,8 @@ def _free_full_params(self, params: Optional[List[Parameter]] = None) -> None:
         """Free up storage for full parameters."""
         if params is None:
             params = self.params
+        self.is_not_first_batch = False
+        self.is_full_params_first_rebuilt = True
         self.has_full_params = False
 
         current_stream = torch.cuda.current_stream()
@@ -2279,8 +2411,11 @@ def local_metadata_dict(self) -> Dict[str, Any]:
                         backing_param_name = m.module.flat_param_names[i]
                         names, shapes, numels = m.module.metadata(i)
                     else:
+                        # TODO(shikaili): Understand this change.
                         assert len(m._param_name_groups[i]) == 1
-                        backing_param_name = m._param_name_groups[i][0]
+                        backing_param_name = m._param_name_groups[
+                            m._num_flatten_params
+                        ][i - m._num_flatten_params]
                         names = [backing_param_name]
                         shapes = [p._orig_size]
                         numels = [p._orig_size.numel()]
@@ -2396,12 +2531,50 @@ def _cast_fp32_param_shards_to_fp16(self, params: Optional[List[Parameter]] = No
             for p in params:
                 assert p._fp16_shard is not None
                 alloc_storage_(p._fp16_shard, size=p._fp32_shard.size())
-                p._fp16_shard.copy_(
-                    # If move_params_to_cpu is True, this will be non-blocking
-                    # because _fp32_shard is pinned, otherwise it's a no-op.
-                    p._fp32_shard.to(p._fp16_shard.device, non_blocking=True)
-                )
-                p.data = p._fp16_shard
+                if self._is_fp8_compute and _is_fp8_dtype(p._fp16_shard.dtype):
+                    assert isinstance(p, FlatParameter)
+                    assert len(p._param_infos) == len(p._param_numels)
+                    numel_per_shard = p.numel()
+                    offset = -numel_per_shard * self.rank
+                    for i in range(len(p._param_infos)):
+                        _, m, n = p._param_infos[i]
+                        numel = p._param_numels[i]
+                        if offset + numel <= 0 or offset >= numel_per_shard:
+                            offset += numel
+                            continue
+                        fp8_dtype_forward = te.fp8.get_fp8_te_dtype(
+                            m.fp8_meta["recipe"], fprop_tensor=True
+                        )
+                        if not m.fp8_initialized:
+                            m.fp8_init(
+                                num_gemms=2 if isinstance(m, te.LayerNormMLP) else 1
+                            )
+                        begin = max(offset, 0)
+                        end = min(offset + numel, numel_per_shard)
+                        cast_to_fp8(
+                            p._fp32_shard[begin:end].bfloat16(),
+                            m.fp8_meta["scaling_fwd"],
+                            (
+                                FP8FwdTensors.GEMM2_WEIGHT
+                                if n == "fc2_weight"
+                                else FP8FwdTensors.GEMM1_WEIGHT
+                            ),
+                            fp8_dtype_forward,
+                            out=p._fp16_shard[begin:end],
+                        )
+                        offset += numel
+                    p.data = p._fp16_shard.view(
+                            torch.float8_e4m3fn
+                            if fp8_dtype_forward == DType.kFloat8E4M3
+                            else torch.float8_e5m2
+                    )
+                else:
+                    p._fp16_shard.copy_(
+                        # If move_params_to_cpu is True, this will be non-blocking
+                        # because _fp32_shard is pinned, otherwise it's a no-op.
+                        p._fp32_shard.to(p._fp16_shard.device, non_blocking=True)
+                    )
+                    p.data = p._fp16_shard
         torch.cuda.current_stream().wait_stream(self._streams["fp32_to_fp16"])
 
     @torch.no_grad()
diff --git a/fairscale/nn/misc/__init__.py b/fairscale/nn/misc/__init__.py
index 71a34cae3..44999f0ca 100644
--- a/fairscale/nn/misc/__init__.py
+++ b/fairscale/nn/misc/__init__.py
@@ -9,7 +9,7 @@
 # in favor of fairscale.nn.checkpoint.checkpoint_wrapper.
 from fairscale.nn.checkpoint import checkpoint_wrapper
 
-from .flatten_params_wrapper import FlattenParamsWrapper
+from .flatten_params_wrapper import FlatParameter, FlattenParamsWrapper
 from .param_bucket import GradBucket, ParamBucket
 
 __all__: List[str] = []
diff --git a/fairscale/nn/misc/flatten_params_wrapper.py b/fairscale/nn/misc/flatten_params_wrapper.py
index 38265dd2b..ae2f3d792 100644
--- a/fairscale/nn/misc/flatten_params_wrapper.py
+++ b/fairscale/nn/misc/flatten_params_wrapper.py
@@ -486,7 +486,8 @@ def load_state_dict(
             return super().load_state_dict(state_dict, strict)
 
     def forward(self, *inputs: Any, **kwinputs: Any) -> Any:
-        self._unflatten_params_as_views()
+        if getattr(self, "is_first_batch", False):
+            self._unflatten_params_as_views()
         return self.module(*inputs, **kwinputs)
 
     def get_param_views(self, external_data_list: Optional[List[Optional[Tensor]]] = None) -> Iterator[Tensor]:

From fa9cf77e947cf551443f873585525288c3023e42 Mon Sep 17 00:00:00 2001
From: Shikai Li <shikaili@fb.com>
Date: Tue, 26 Mar 2024 21:28:16 -0700
Subject: [PATCH 13/20] Debug non-determinism issues.

This commit works with a 4 GPU run on SMALL model with FSDP and PP
enabled.
---
 .../fully_sharded_data_parallel.py            | 166 +++++++++---------
 fairscale/nn/misc/__init__.py                 |   2 +-
 2 files changed, 80 insertions(+), 88 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index ca744cafe..e97f89ab9 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -48,7 +48,8 @@
 from transformer_engine.pytorch.cpp_extensions import cast_to_fp8, DType, FP8FwdTensors
 from transformer_engine.pytorch.fp8 import amax_and_scale_update, FP8GlobalStateManager
 
-from fairscale.nn.misc import FlatParameter, FlattenParamsWrapper
+from fairscale.nn.misc import FlattenParamsWrapper
+from fairscale.nn.misc.flatten_params_wrapper import FlatParameter
 from fairscale.nn.wrap import auto_wrap, config_auto_wrap_policy, enable_wrap
 from fairscale.utils.containers import apply_to_tensors
 from fairscale.utils.parallel import (
@@ -1273,7 +1274,7 @@ def _lazy_init(self) -> None:
             # ``optim.step()`` is done before we all-gather parameters.
             self._wait_for_previous_optim_step()
 
-    def _shard_dtype(self, p: Parameter) -> torch.dtype:
+    def _param_dtype(self, p: Parameter) -> torch.dtype:
         """
         Return the dtype to use for the sharded parameters.
 
@@ -1346,7 +1347,7 @@ def _init_param_attributes(self, p: Parameter) -> None:
             # from _fp32_shard) as needed. If offloading params to CPU, the
             # dtype of the fp16 shard will depend on the *`compute_dtype`*.
             p._fp16_shard = torch.zeros_like(
-                p._fp32_shard, device=self.compute_device, dtype=self._shard_dtype(p)
+                p._fp32_shard, device=self.compute_device, dtype=self._param_dtype(p)
             )
             free_storage_(p._fp16_shard)
 
@@ -1368,7 +1369,7 @@ def _init_param_attributes(self, p: Parameter) -> None:
             p._full_param_padded = torch.zeros(
                 p.data.numel() * self.world_size,
                 device=self.compute_device,
-                dtype=self._shard_dtype(p),
+                dtype=self._param_dtype(p),
             )
             free_storage_(p._full_param_padded)
 
@@ -1780,11 +1781,19 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
             if hasattr(param._linked_param, "_is_shared") and param._linked_param._is_shared:
                 param = param._linked_param
 
-        if param.grad is None:
+        # Prefer to use `param.main_grad` with higher precision in reduction than
+        # `param.grad` with equal or lower precision.
+        grad = param.grad
+        main_grad = getattr(param, "main_grad", None)
+        to_reduce_grad = main_grad if (main_grad is not None and not param.main_grad.eq(0.0).all()) else grad
+
+        if to_reduce_grad is None:
             return
 
-        if param.grad.requires_grad:
-            raise RuntimeError("FSDP only works with gradients that don't require gradients")
+        if to_reduce_grad.requires_grad:
+            raise RuntimeError(
+                    "FSDP only works with gradients that don't require gradients"
+                    )
 
         if self._should_free_in_backward():
             # Free full params.
@@ -1799,49 +1808,52 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
         # Switch to FP32 shard after backward.
         self._use_fp32_param_shard([param])
 
+        # Accumulate gradients manually if in FP32 instead of default precision.
         if self.fp32_reduce_scatter:
-            if getattr(param, "unsharded_main_grad", None) is None:
-                param.unsharded_main_grad = param.grad.to(torch.float32)
-            else:
-                param.unsharded_main_grad.add_(param.grad.data)
-            param.grad = None
+            # If `param.grad` is `None`, then fp32 reduction is already happening at
+            # `param.main_grad`, nothing need to be done here.
+            if param.grad is not None:
+                if getattr(param, "unsharded_main_grad", None) is None:
+                    param.unsharded_main_grad = param.grad.to(torch.float32)
+                else:
+                    param.unsharded_main_grad.add_(param.grad.to(torch.float32))
+                # Resets `param.grad` to avoid PyTorch accumulation.
+                param.grad = None
 
         if not self._require_backward_grad_sync:
             return
 
-
         # Wait for all work in the current stream to finish, then start the
         # reductions in post_backward stream.
         self._streams["post_backward"].wait_stream(torch.cuda.current_stream())
         with torch.cuda.stream(self._streams["post_backward"]):
 
-            if self.fp32_reduce_scatter:
-                # Cast grad to FP32.
-                orig_grad_data = param.unsharded_main_grad.data
-            else:
-                orig_grad_data = param.grad.data
+            # Prefer to use "unsharded_main_grad" with higher precision if exsits.
+            unsharded_main_grad = getattr(param, "unsharded_main_grad", None)
+            to_reduce_grad = unsharded_main_grad if unsharded_main_grad is not None else to_reduce_grad
+
+            # Clear grad on the tensor, so any repeated gradient computations do not interfere with this reduction.
+            # 1. For sharded parameters, we will asynchronously accumulate the reduced gradient into
+            # `param._saved_grad_shard` which will be re-sharded to `param.grad`/`param.main_grad` later after
+            # finalization.
+            # 2. For unsharded parameters, we will directly asynchronously accumulate the reduce gradient into
+            # `param.grad`/`param.main_grad`.
+            param.grad = None
+            param.main_grad = None
+            # `param.unsharded_main_grad` is no longer usefual and will be created again inside
+            # `_post_backward_hook` in the next first microbatch.
+            param.unsharded_main_grad = None
 
             if self.gradient_predivide_factor > 1:
                 # Average grad by world_size for consistency with PyTorch DDP.
-                if getattr(param, "unsharded_main_grad", None) is not None:
-                    param.unsharded_main_grad.data.div_(self.gradient_predivide_factor)
-                else:
-                    param.grad.data.div_(self.gradient_predivide_factor)
+                to_reduce_grad.data.div_(self.gradient_predivide_factor)
 
             if param._is_sharded:
                 assert self._reducer is not None
-                # Save the unsharded grad for reduction. We will asynchronously accumulate the reduced gradient into
-                # param._saved_grad_shard. If this FSDP module was called multiple times it's possible that multiple
-                # gradient reductions will happen in an undefined order. But addition commutes, so this order doesn't
-                # matter, neglecting rounding.
-                if getattr(param, "unsharded_main_grad", None) is not None:
-                    grad = param.unsharded_main_grad.data
-                    param.unsharded_main_grad = None
-                else:
-                    grad = param.grad.data
-                    param.grad = None
 
-                # Clear grad on the tensor, so any repeated gradient computations do not interfere with this reduction.
+                # If this FSDP module was called multiple times it's possible that multiple gradient reductions will
+                # happen in an undefined order. But addition commutes, so this order doesn't matter, neglecting
+                # rounding.
                 #
                 # The effect on memory consumption is not usually significant. No extra memory is allocated if this
                 # module is called only once, reduction happens quickly, or the tensor is bucketed. If the module is
@@ -1854,37 +1866,29 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                 # unsharded gradients allocated; one for a pending reduction, and one for gradient computation.
                 callback_fn = functools.partial(self._post_reduction_hook, param)
                 self._reducer.reduce_scatter_async(
-                    grad, group=self.process_group_reduce_scatter, callback_fn=callback_fn
+                    to_reduce_grad,
+                    group=self.process_group_reduce_scatter,
+                    callback_fn=callback_fn,
                 )
-            elif self.fp8_allgather:
-                if param.grad is not None:
-                    if self.fp32_reduce_scatter:
-                        param.unsharded_main_grad = param.grad.to(torch.float32)
-                    else:
-                        param.unsharded_main_grad = param.grad
-                    param.grad = None
-
+            else:
+                # Unsharded parameters only happens for word_size == 1 or fp8_allgather
+                assert self.world_size == 1 or (self._is_fp8_compute and not isinstance(param, FlatParameter))
+                if self.world_size > 1:
                     torch.distributed.all_reduce(
-                        param.unsharded_main_grad,
+                        to_reduce_grad,
                         group=self.process_group_reduce_scatter,
                     )
-                    self._post_reduction_hook(param, param.unsharded_main_grad)
-            else:
-                # Currently the only way for _is_sharded to be False is if
-                # world_size == 1. This could be relaxed in the future, in which
-                # case grads should be all-reduced here.
-                assert self.world_size == 1
-                if getattr(param, "unsharded_main_grad", None) is not None:
-                    self._post_reduction_hook(param, param.unsharded_main_grad)
-                else:
-                    self._post_reduction_hook(param, param.grad)
+                self._post_reduction_hook(param, to_reduce_grad)
 
             # After _post_backward_hook returns, orig_grad_data will eventually
             # go out of scope, at which point it could otherwise be freed for
             # further reuse by the main stream while the div/reduce_scatter/copy
             # are underway in the post_backward stream. See:
             # github.com/NVIDIA/apex/blob/master/apex/parallel/distributed.py
-            orig_grad_data.record_stream(self._streams["post_backward"])
+            for g in (grad, main_grad, unsharded_main_grad):
+                if g is not None:
+                    g.data.record_stream(self._streams["post_backward"])
+
             if self._reduce_scatter_free_event_queue is not None:
                 release_full_grad_event = torch.cuda.Event()
                 release_full_grad_event.record()
@@ -1895,17 +1899,12 @@ def _post_reduction_hook(self, param: Parameter, reduced_grad: torch.Tensor) ->
         """Hook to call on each param after the reduce-scatter."""
         assert torch.cuda.current_stream() == self._streams["post_backward"]
         self.assert_state(TrainingState.BACKWARD_POST)
+
+        assert not (self.fp32_reduce_scatter and reduced_grad.dtype != param.dtype)
+
         if self.gradient_postdivide_factor > 1:
             # Average grad by world_size for consistency with PyTorch DDP.
             reduced_grad.data.div_(self.gradient_postdivide_factor)
-        # Cast grad to param's dtype (typically FP32). Note: we do this
-        # before the move_grads_to_cpu step so that this entire hook remains
-        # non-blocking. The downside is a bit more D2H transfer in that case.
-        if self.fp32_reduce_scatter:
-            orig_param_grad_data = reduced_grad.data
-            # reduced_grad.data = reduced_grad.data.to(dtype=param.data.dtype)
-            # Don't let this memory get reused until after the transfer.
-            orig_param_grad_data.record_stream(torch.cuda.current_stream())
 
         if param._is_sharded:
             # Accumulate into the gradient shard.
@@ -1918,16 +1917,9 @@ def _post_reduction_hook(self, param: Parameter, reduced_grad: torch.Tensor) ->
                 param._saved_grad_shard.data += reduced_grad.data
             reduced_grad = param._saved_grad_shard.data
 
-            if getattr(param, "unsharded_main_grad", None) is not None:
-                free_storage_(param.unsharded_main_grad.data)
-                param.unsharded_main_grad = None
-
         elif param.grad is None:
-        # TODO(shikaili): Still having missing grad/main_grad issue.
             if self.fp32_reduce_scatter:
                 param.main_grad = reduced_grad.data
-            else:
-                param.grad = reduced_grad.data
 
         # Optionally move gradients to CPU, typically used if one is running the optimizer on the CPU. Once the full
         # backwards pass completes, we will set `.grad` to the CPU copy.
@@ -2147,6 +2139,7 @@ def _rebuild_full_params(
                     d = {info[0]: info[1] for info in p._param_infos}
                     for n, m in d.items():
                         # Previous iteration was grad_enabled
+                        # assert m.fp8_initialized
                         if not m.fp8_initialized:
                             m.fp8_init(
                                 num_gemms=2 if isinstance(m, te.LayerNormMLP) else 1
@@ -2156,22 +2149,22 @@ def _rebuild_full_params(
                                 FP8GlobalStateManager.copy_amax_from_global_buffer(
                                     m.fp8_meta, forward=True
                                 )
-                            amax_and_scale_update(
-                                m.fp8_meta,
-                                True,
-                                update_weight_scale_inv=self.is_full_params_first_rebuilt,
-                            )
-                            if not_from_recursive:
-                                FP8GlobalStateManager.set_amax_buffer_key_deletion(
-                                    m.fp8_meta, forward=True
+                                amax_and_scale_update(
+                                    m.fp8_meta,
+                                    True,
+                                    update_weight_scale_inv=self.is_full_params_first_rebuilt,
                                 )
-                        else:
-                            amax_and_scale_update(
-                                m.fp8_meta,
-                                True,
-                                update_weight_scale_inv=self.is_full_params_first_rebuilt,
-                            )
-                torch.cuda.current_stream().wait_stream(self._streams["fp32_to_fp16"])
+                                if not_from_recursive:
+                                    FP8GlobalStateManager.set_amax_buffer_key_deletion(
+                                        m.fp8_meta, forward=True
+                                    )
+                            else:
+                                amax_and_scale_update(
+                                    m.fp8_meta,
+                                    True,
+                                    update_weight_scale_inv=self.is_full_params_first_rebuilt,
+                                )
+            torch.cuda.current_stream().wait_stream(self._streams["fp32_to_fp16"])
 
         output_tensors: List[Tuple[torch.Tensor, bool]] = []
 
@@ -2338,7 +2331,7 @@ def _prep_grads_for_backward(self) -> None:
                 _is_te_module_with_weights(info[1]) for info in p._param_infos
             ):
                 if getattr(p, "main_grad", None) is None:
-                    p.main_grad = torch.empty_like(p, dtype=torch.float)
+                    p.main_grad = torch.empty_like(p, dtype=torch.float32)
                 main_grad_views = p.get_param_views(p.main_grad)
                 for (_, m, n), main_grad in zip(p._param_infos, main_grad_views):
                     getattr(m, n).main_grad = main_grad
@@ -2411,8 +2404,6 @@ def local_metadata_dict(self) -> Dict[str, Any]:
                         backing_param_name = m.module.flat_param_names[i]
                         names, shapes, numels = m.module.metadata(i)
                     else:
-                        # TODO(shikaili): Understand this change.
-                        assert len(m._param_name_groups[i]) == 1
                         backing_param_name = m._param_name_groups[
                             m._num_flatten_params
                         ][i - m._num_flatten_params]
@@ -2542,6 +2533,7 @@ def _cast_fp32_param_shards_to_fp16(self, params: Optional[List[Parameter]] = No
                         if offset + numel <= 0 or offset >= numel_per_shard:
                             offset += numel
                             continue
+                        assert _is_te_module_with_weights(m)
                         fp8_dtype_forward = te.fp8.get_fp8_te_dtype(
                             m.fp8_meta["recipe"], fprop_tensor=True
                         )
diff --git a/fairscale/nn/misc/__init__.py b/fairscale/nn/misc/__init__.py
index 44999f0ca..71a34cae3 100644
--- a/fairscale/nn/misc/__init__.py
+++ b/fairscale/nn/misc/__init__.py
@@ -9,7 +9,7 @@
 # in favor of fairscale.nn.checkpoint.checkpoint_wrapper.
 from fairscale.nn.checkpoint import checkpoint_wrapper
 
-from .flatten_params_wrapper import FlatParameter, FlattenParamsWrapper
+from .flatten_params_wrapper import FlattenParamsWrapper
 from .param_bucket import GradBucket, ParamBucket
 
 __all__: List[str] = []

From 6fa19e07dcff85e7ca4c42c3a01b99016c284061 Mon Sep 17 00:00:00 2001
From: Shikai Li <shikaili@fb.com>
Date: Thu, 28 Mar 2024 16:22:18 -0700
Subject: [PATCH 14/20] Moves amax update logic into params downcasting
 function.

---
 .../fully_sharded_data_parallel.py            | 92 +++++++++----------
 1 file changed, 44 insertions(+), 48 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index e97f89ab9..c4710cdf3 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -2130,42 +2130,6 @@ def _rebuild_full_params(
             caller to free the full-sized param. This will be ``None`` if
             ``force_full_precision=False`` and the full params are already gathered.
         """
-        if self._is_fp8_compute:
-            # Need to use fp32_to_fp16 stream since _cast_fp32_param_shards_to_fp16 depends on this block.
-            with torch.no_grad(), torch.cuda.stream(self._streams["fp32_to_fp16"]):
-                for p in self.params:
-                    if not isinstance(p, FlatParameter):
-                        continue
-                    d = {info[0]: info[1] for info in p._param_infos}
-                    for n, m in d.items():
-                        # Previous iteration was grad_enabled
-                        # assert m.fp8_initialized
-                        if not m.fp8_initialized:
-                            m.fp8_init(
-                                num_gemms=2 if isinstance(m, te.LayerNormMLP) else 1
-                            )
-                        if m.fp8_meta.get("update_amax_and_scale_fwd", False):
-                            if m.fp8_meta["recipe"].reduce_amax:
-                                FP8GlobalStateManager.copy_amax_from_global_buffer(
-                                    m.fp8_meta, forward=True
-                                )
-                                amax_and_scale_update(
-                                    m.fp8_meta,
-                                    True,
-                                    update_weight_scale_inv=self.is_full_params_first_rebuilt,
-                                )
-                                if not_from_recursive:
-                                    FP8GlobalStateManager.set_amax_buffer_key_deletion(
-                                        m.fp8_meta, forward=True
-                                    )
-                            else:
-                                amax_and_scale_update(
-                                    m.fp8_meta,
-                                    True,
-                                    update_weight_scale_inv=self.is_full_params_first_rebuilt,
-                                )
-            torch.cuda.current_stream().wait_stream(self._streams["fp32_to_fp16"])
-
         output_tensors: List[Tuple[torch.Tensor, bool]] = []
 
         def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
@@ -2176,7 +2140,6 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
                 custom_output_tensor (torch.Tensor, Optional): if not None, this
                 tensor contains the data we just gathered.
             """
-            p_fp16_shard_size = -1
             if custom_output_tensor is not None:
                 assert p._is_sharded
                 p.data = custom_output_tensor
@@ -2185,7 +2148,6 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
                 if (self.mixed_precision or self.move_params_to_cpu) and not force_full_precision:
                     assert p._fp16_shard is not None
                     p.data = p._fp16_shard
-                    p_fp16_shard_size = p._fp16_shard.storage().size()
                     output_tensors.append((p.data, True))
                 else:
                     # Here p.data == p._fp32_shard, so it's not safe to free.
@@ -2234,8 +2196,10 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
                 event.synchronize()
 
         with torch.cuda.stream(self._streams["all_gather"]):
-            if (self.mixed_precision or self.move_params_to_cpu) and not force_full_precision:
-                self._cast_fp32_param_shards_to_fp16()
+            if (
+                self.mixed_precision or self.move_params_to_cpu
+            ) and not force_full_precision:
+                self._cast_fp32_param_shards_to_fp16(not_from_recursive=not_from_recursive)
 
             if self.move_params_to_cpu:
                 if force_full_precision:
@@ -2243,7 +2207,7 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
                     # use pinned memory. Otherwise move p.data to the compute
                     # device.
                     if self.params[0].dtype == self.compute_dtype:
-                        self._cast_fp32_param_shards_to_fp16()
+                        self._cast_fp32_param_shards_to_fp16(not_from_recursive=not_from_recursive)
                     else:
                         for p in self.params:
                             p.data = p.data.to(self.compute_device)
@@ -2514,33 +2478,64 @@ def _use_fp32_param_shard(self, params: Optional[List[Parameter]] = None) -> Non
             p.data = p._fp32_shard
 
     @torch.no_grad()
-    def _cast_fp32_param_shards_to_fp16(self, params: Optional[List[Parameter]] = None) -> None:
+    def _cast_fp32_param_shards_to_fp16(
+        self, params: Optional[List[Parameter]] = None,
+        not_from_recursive: bool = False,
+    ) -> None:
         """Cast FP32 param shard to FP16 for a list of params."""
         if params is None:
             params = self.params
+
         with torch.cuda.stream(self._streams["fp32_to_fp16"]):
             for p in params:
                 assert p._fp16_shard is not None
                 alloc_storage_(p._fp16_shard, size=p._fp32_shard.size())
+
                 if self._is_fp8_compute and _is_fp8_dtype(p._fp16_shard.dtype):
                     assert isinstance(p, FlatParameter)
                     assert len(p._param_infos) == len(p._param_numels)
+
                     numel_per_shard = p.numel()
                     offset = -numel_per_shard * self.rank
                     for i in range(len(p._param_infos)):
                         _, m, n = p._param_infos[i]
+                        assert _is_te_module_with_weights(m)
+
+                        if not m.fp8_initialized:
+                            m.fp8_init(
+                                num_gemms=2 if isinstance(m, te.LayerNormMLP) else 1
+                            )
+
+                        if self.is_full_params_first_rebuilt:
+                            if m.fp8_meta.get("update_amax_and_scale_fwd", False):
+                                if m.fp8_meta["recipe"].reduce_amax:
+                                    FP8GlobalStateManager.copy_amax_from_global_buffer(
+                                        m.fp8_meta, forward=True
+                                    )
+                                    amax_and_scale_update(
+                                        m.fp8_meta,
+                                        True,
+                                        update_weight_scale_inv=True,
+                                    )
+                                    if not_from_recursive:
+                                        FP8GlobalStateManager.set_amax_buffer_key_deletion(
+                                            m.fp8_meta, forward=True
+                                        )
+                                else:
+                                    amax_and_scale_update(
+                                        m.fp8_meta,
+                                        True,
+                                        update_weight_scale_inv=True,
+                                    )
+
                         numel = p._param_numels[i]
                         if offset + numel <= 0 or offset >= numel_per_shard:
                             offset += numel
                             continue
-                        assert _is_te_module_with_weights(m)
+
                         fp8_dtype_forward = te.fp8.get_fp8_te_dtype(
                             m.fp8_meta["recipe"], fprop_tensor=True
                         )
-                        if not m.fp8_initialized:
-                            m.fp8_init(
-                                num_gemms=2 if isinstance(m, te.LayerNormMLP) else 1
-                            )
                         begin = max(offset, 0)
                         end = min(offset + numel, numel_per_shard)
                         cast_to_fp8(
@@ -2567,7 +2562,8 @@ def _cast_fp32_param_shards_to_fp16(self, params: Optional[List[Parameter]] = No
                         p._fp32_shard.to(p._fp16_shard.device, non_blocking=True)
                     )
                     p.data = p._fp16_shard
-        torch.cuda.current_stream().wait_stream(self._streams["fp32_to_fp16"])
+
+        self._streams["all_gather"].wait_stream(self._streams["fp32_to_fp16"])
 
     @torch.no_grad()
     def _free_fp16_param_shard(self, params: Optional[List[Parameter]] = None) -> None:

From 80ffd546e6a0bb8da7efd57b123a075ed95c0144 Mon Sep 17 00:00:00 2001
From: Shikai Li <shikaili@fb.com>
Date: Thu, 28 Mar 2024 17:13:57 -0700
Subject: [PATCH 15/20] Cleans up code.

---
 .../fully_sharded_data_parallel.py            | 67 +++++++++----------
 1 file changed, 33 insertions(+), 34 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index c4710cdf3..986ee8603 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -269,6 +269,9 @@ class FullyShardedDataParallel(nn.Module):
         fp32_reduce_scatter (bool, Optional):
             if ``True``, then reduce-scatter gradients in FP32. This is only
             relevant when *``mixed_precision``* is ``True``.
+        fp8_all_gather (bool, Optional):
+            if ``True``, then all-gather weights/gradients in FP8. This is only
+            relevant when *``mixed_precision``* is ``True``.
         flatten_parameters (bool, Optional):
             if ``True``, flatten parameters into a single contiguous tensor,
             which improves training speed.
@@ -364,6 +367,7 @@ def __init__(
         disable_reshard_on_root: bool = True,
         mixed_precision: bool = False,
         fp32_reduce_scatter: bool = False,
+        fp8_all_gather: bool = False,
         flatten_parameters: bool = True,
         move_params_to_cpu: bool = False,
         compute_dtype: Optional[torch.dtype] = None,
@@ -384,7 +388,6 @@ def __init__(
         limit_reduce_scatter_events: bool = False,
         cast_input: bool = True,
         should_validate_process_group: bool = True,
-        fp8_allgather: bool = False,
     ):
         try:
             import torch._C
@@ -437,6 +440,7 @@ def __init__(
         self.mixed_precision = mixed_precision
         self.cast_input = cast_input
         self.fp32_reduce_scatter = fp32_reduce_scatter
+        self.fp8_all_gather = fp8_all_gather
         self.flatten_parameters = flatten_parameters
         self.move_params_to_cpu = move_params_to_cpu or cpu_offload
         self.compute_dtype = compute_dtype or (torch.float16 if mixed_precision else torch.float32)
@@ -451,7 +455,6 @@ def __init__(
         self.force_input_to_fp32 = force_input_to_fp32
         self.verbose = verbose
         self.state_dict_on_rank_0_only = state_dict_on_rank_0_only
-        self.fp8_allgather = fp8_allgather
         # Experimental feature for now. Use at your own risk.
         self.ssd_offload = True if offload_config and offload_config.offload_type == "ssd_offload" else False
 
@@ -589,6 +592,7 @@ def __init__(
         self.dont_wait_current_stream_for_post_all_gather = False
         self._all_gather_free_event_queue = _FreeEventQueue() if limit_all_gather_events else None
         self._reduce_scatter_free_event_queue = _FreeEventQueue() if limit_reduce_scatter_events else None
+        self.is_full_params_first_rebuilt = True
 
     @property
     def _is_fp8_compute(self) -> bool:
@@ -1281,9 +1285,8 @@ def _param_dtype(self, p: Parameter) -> torch.dtype:
         Returns:
             The dtype to use for the sharded parameters.
         """
-        if self._is_fp8_compute and not isinstance(p, FlatParameter):
-            # Assume non flattened are precision critical like norm
-            assert not p._is_sharded
+        if self._is_fp8_compute and (not self.fp8_all_gather or
+                                     not isinstance(p, FlatParameter)):
             return torch.bfloat16
         else:
             return self.compute_dtype
@@ -1300,13 +1303,13 @@ def _init_param_attributes(self, p: Parameter) -> None:
             ``_orig_size``: the size of the original Parameter (before sharding)
 
         The remaining attributes are set here:
-            ``_fp32_shard``: a single shard of the parameters in full precision
-                (typically FP32, but this is dependent on the dtype of the model
-                as it's passed in by the user). This can be on CPU or GPU
+            ``_fp32_shard``: This will be a single shard of the parameters in
+                full precision (typically FP32, but this is dependent on the dtype of
+                the model as it's passed in by the user). This can be on CPU or GPU
                 depending on the value of *``move_params_to_cpu``*.
-            ``_fp16_shard``: This will be a single shard of the parameters in FP16, used for all-gather.
-                This can be in FP16 or FP32 depending on the value of *``compute_dtype``* and
-                if params are offloaded to CPU.
+            ``_fp16_shard``: This will be a single shard of the parameters
+                used for all-gather. This can be in FP8, FP16 or FP32 depending on the value
+                of *``compute_dtype``*, *``fp8_all_gather``*, *``move_params_to_cpu``*..
             ``_full_param_padded``: the full weight (padded to be evenly
                 divisible by ``world_size``), used for computation in the
                 forward and backward pass. This will be resized in place and
@@ -1326,6 +1329,7 @@ def _init_param_attributes(self, p: Parameter) -> None:
 
         if self.mixed_precision:
             assert p._fp32_shard.dtype == torch.float32
+
         if self.move_params_to_cpu:
             assert p._fp32_shard.device == torch.device("cpu")
 
@@ -1339,7 +1343,6 @@ def _init_param_attributes(self, p: Parameter) -> None:
                 p.data = p._fp32_shard
 
         if self.move_params_to_cpu or self.mixed_precision:
-
             # In mixed precision mode, we maintain a reduced precision
             # (typically FP16) parameter shard on compute_device for performing
             # the computation in the forward/backward pass. We resize the
@@ -1350,11 +1353,7 @@ def _init_param_attributes(self, p: Parameter) -> None:
                 p._fp32_shard, device=self.compute_device, dtype=self._param_dtype(p)
             )
             free_storage_(p._fp16_shard)
-
-        if self.mixed_precision:
-            assert p._fp32_shard.dtype == torch.float32
-
-        if not self.mixed_precision and not self.move_params_to_cpu:
+        else:
             # use _fp32_shard if you are not in using mixed precision or
             # offloading params and grads to CPU.
             p._fp16_shard = None
@@ -1435,8 +1434,8 @@ def _setup_streams(self) -> None:
             return
 
         if torch.cuda.is_available():
-            # Stream to move main FP32 params (may be on CPU) to FP16 for forward.
-            self._streams["fp32_to_fp16"] = torch.cuda.Stream()
+            # Stream to move main FP32 params (may be on CPU) to FP32/FP16/FP8 for forward.
+            self._streams["cast_param"] = torch.cuda.Stream()
             # Stream for all-gathering parameters.
             self._streams["all_gather"] = torch.cuda.Stream()
             # Stream for overlapping grad reduction with the backward pass.
@@ -1472,7 +1471,7 @@ def _wait_for_previous_optim_step(self) -> None:
         if not torch.cuda.is_available():
             return
         if self.mixed_precision or self.move_params_to_cpu:
-            self._streams["fp32_to_fp16"].wait_stream(torch.cuda.current_stream())
+            self._streams["cast_param"].wait_stream(torch.cuda.current_stream())
         else:
             self._streams["all_gather"].wait_stream(torch.cuda.current_stream())
 
@@ -1799,7 +1798,7 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
             # Free full params.
             self._free_full_params([param])
 
-        if self.mixed_precision and (self._require_backward_grad_sync or self.reshard_after_forward) and not self.fp8_allgather:
+        if self.mixed_precision and (self._require_backward_grad_sync or self.reshard_after_forward) and not self.fp8_all_gather:
             # This is a no-op if reshard_after_forward is True, since we already
             # free the param shard when rebuilding the full params in the
             # pre_backward_hook.
@@ -1816,7 +1815,7 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                 if getattr(param, "unsharded_main_grad", None) is None:
                     param.unsharded_main_grad = param.grad.to(torch.float32)
                 else:
-                    param.unsharded_main_grad.add_(param.grad.to(torch.float32))
+                    param.unsharded_main_grad.data.add_(param.grad.to(torch.float32))
                 # Resets `param.grad` to avoid PyTorch accumulation.
                 param.grad = None
 
@@ -1871,7 +1870,7 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                     callback_fn=callback_fn,
                 )
             else:
-                # Unsharded parameters only happens for word_size == 1 or fp8_allgather
+                # Unsharded parameters only happens for word_size == 1 or fp8_all_gather
                 assert self.world_size == 1 or (self._is_fp8_compute and not isinstance(param, FlatParameter))
                 if self.world_size > 1:
                     torch.distributed.all_reduce(
@@ -1885,9 +1884,7 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
             # further reuse by the main stream while the div/reduce_scatter/copy
             # are underway in the post_backward stream. See:
             # github.com/NVIDIA/apex/blob/master/apex/parallel/distributed.py
-            for g in (grad, main_grad, unsharded_main_grad):
-                if g is not None:
-                    g.data.record_stream(self._streams["post_backward"])
+            to_reduce_grad.data.record_stream(self._streams["post_backward"])
 
             if self._reduce_scatter_free_event_queue is not None:
                 release_full_grad_event = torch.cuda.Event()
@@ -1900,7 +1897,7 @@ def _post_reduction_hook(self, param: Parameter, reduced_grad: torch.Tensor) ->
         assert torch.cuda.current_stream() == self._streams["post_backward"]
         self.assert_state(TrainingState.BACKWARD_POST)
 
-        assert not (self.fp32_reduce_scatter and reduced_grad.dtype != param.dtype)
+        # assert not (self.fp32_reduce_scatter and reduced_grad.dtype != param.dtype)
 
         if self.gradient_postdivide_factor > 1:
             # Average grad by world_size for consistency with PyTorch DDP.
@@ -1914,12 +1911,14 @@ def _post_reduction_hook(self, param: Parameter, reduced_grad: torch.Tensor) ->
                 assert (
                     param._saved_grad_shard.shape == reduced_grad.shape
                 ), f"{param._saved_grad_shard.shape} vs {reduced_grad.shape}"
-                param._saved_grad_shard.data += reduced_grad.data
+                param._saved_grad_shard.data.add_(reduced_grad.data)
             reduced_grad = param._saved_grad_shard.data
 
         elif param.grad is None:
             if self.fp32_reduce_scatter:
                 param.main_grad = reduced_grad.data
+            else:
+                param.grad = reduced_grad.data
 
         # Optionally move gradients to CPU, typically used if one is running the optimizer on the CPU. Once the full
         # backwards pass completes, we will set `.grad` to the CPU copy.
@@ -2199,7 +2198,7 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
             if (
                 self.mixed_precision or self.move_params_to_cpu
             ) and not force_full_precision:
-                self._cast_fp32_param_shards_to_fp16(not_from_recursive=not_from_recursive)
+                self._cast_params_for_all_gather(not_from_recursive=not_from_recursive)
 
             if self.move_params_to_cpu:
                 if force_full_precision:
@@ -2207,7 +2206,7 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
                     # use pinned memory. Otherwise move p.data to the compute
                     # device.
                     if self.params[0].dtype == self.compute_dtype:
-                        self._cast_fp32_param_shards_to_fp16(not_from_recursive=not_from_recursive)
+                        self._cast_params_for_all_gather(not_from_recursive=not_from_recursive)
                     else:
                         for p in self.params:
                             p.data = p.data.to(self.compute_device)
@@ -2478,7 +2477,7 @@ def _use_fp32_param_shard(self, params: Optional[List[Parameter]] = None) -> Non
             p.data = p._fp32_shard
 
     @torch.no_grad()
-    def _cast_fp32_param_shards_to_fp16(
+    def _cast_params_for_all_gather(
         self, params: Optional[List[Parameter]] = None,
         not_from_recursive: bool = False,
     ) -> None:
@@ -2486,7 +2485,7 @@ def _cast_fp32_param_shards_to_fp16(
         if params is None:
             params = self.params
 
-        with torch.cuda.stream(self._streams["fp32_to_fp16"]):
+        with torch.cuda.stream(self._streams["cast_param"]):
             for p in params:
                 assert p._fp16_shard is not None
                 alloc_storage_(p._fp16_shard, size=p._fp32_shard.size())
@@ -2563,7 +2562,7 @@ def _cast_fp32_param_shards_to_fp16(
                     )
                     p.data = p._fp16_shard
 
-        self._streams["all_gather"].wait_stream(self._streams["fp32_to_fp16"])
+        self._streams["all_gather"].wait_stream(self._streams["cast_param"])
 
     @torch.no_grad()
     def _free_fp16_param_shard(self, params: Optional[List[Parameter]] = None) -> None:
@@ -2573,7 +2572,7 @@ def _free_fp16_param_shard(self, params: Optional[List[Parameter]] = None) -> No
         current_stream = torch.cuda.current_stream()
         for p in params:
             if p._fp16_shard is not None:
-                # _fp16_shard is allocated in "fp32_to_fp16" stream, so we can't
+                # _fp16_shard is allocated in "cast_param" stream, so we can't
                 # free it until the work in the current stream completes.
                 p._fp16_shard.record_stream(current_stream)
                 free_storage_(p._fp16_shard)

From afb2ca1725e6c6e5d1f2489f08d4c0457df58724 Mon Sep 17 00:00:00 2001
From: Shikai Li <shikaili@fb.com>
Date: Mon, 1 Apr 2024 21:32:45 -0700
Subject: [PATCH 16/20] Fix `main_grad` attribute checking.

- Clean up flatten and non_flatten parameter generation logic.
- Avoid checking `main_grad` attribute all equal to zeros.
---
 .../fully_sharded_data_parallel.py            | 65 ++++++++++---------
 1 file changed, 33 insertions(+), 32 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index 986ee8603..3e6923042 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -505,26 +505,27 @@ def __init__(
 
         # For now, it is either all flatten or none flatten. This will be extended to
         # multiple flatten groups in my next PR.
-        to_be_flatten_params: List[List[Parameter]] = [[]]
-        non_flatten_params = params
-        param_name_groups = [[n] for n in param_names]
-        if self.flatten_parameters:
-            to_be_flatten_params = [
-                [
-                    params[i]
-                    for i in range(len(params))
-                    if "norm_weight" not in param_names[i]
-                ]
-            ]
-            non_flatten_params = [
-                params[i]
-                for i in range(len(params))
-                if "norm_weight" in param_names[i]
-            ]
-            param_name_groups = [
-                [n for n in param_names if "norm_weight" not in n],
-                [n for n in param_names if "norm_weight" in n],
-            ]
+        def should_flatten(name: str) -> bool:
+            # `*_norm_weights` are numerics-sensitive and cannot be quantized to fp8.
+            return self.flatten_parameters and (not self.fp8_all_gather or "norm_weight" not in name)
+
+        to_be_flatten_params: List[List[Parameter]] = [
+            param
+            for param, name in zip(params, param_names)
+            if should_flatten(name)
+        ]
+        if to_be_flatten_params:
+            to_be_flatten_params = [to_be_flatten_params]
+        non_flatten_params: List[List[Parameter]] = [
+            param
+            for param, name in zip(params, param_names)
+            if not should_flatten(name)
+        ]
+        param_name_groups: List[List[str]] = [
+            [n for n in param_names if should_flatten(n)]
+        ] + [
+            [n] for n in param_names if not should_flatten(n)
+        ]
         del param_names
 
         self._fsdp_wrapped_module: nn.Module = FlattenParamsWrapper(
@@ -1769,7 +1770,6 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
         self.assert_state([TrainingState.BACKWARD_PRE, TrainingState.BACKWARD_POST])
         self.training_state = TrainingState.BACKWARD_POST
 
-
         if hasattr(param, "_linked_param"):
             # This links to a shared param. We should finalize the linked param here.
             assert param.shape == (1,), param.shape
@@ -1780,11 +1780,13 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
             if hasattr(param._linked_param, "_is_shared") and param._linked_param._is_shared:
                 param = param._linked_param
 
-        # Prefer to use `param.main_grad` with higher precision in reduction than
-        # `param.grad` with equal or lower precision.
         grad = param.grad
         main_grad = getattr(param, "main_grad", None)
-        to_reduce_grad = main_grad if (main_grad is not None and not param.main_grad.eq(0.0).all()) else grad
+        # Only one of `grad` or `main_grad` can exists. Whenever `main_grad is used for accumulation,
+        # grad should be set as `None`.
+        assert not (grad is not None and main_grad is not None)
+        # Use `grad` or `main_grad` that is not None and avoid invoking a kernel to check all zeros.
+        to_reduce_grad = grad if grad is not None else main_grad
 
         if to_reduce_grad is None:
             return
@@ -1870,8 +1872,8 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                     callback_fn=callback_fn,
                 )
             else:
-                # Unsharded parameters only happens for word_size == 1 or fp8_all_gather
-                assert self.world_size == 1 or (self._is_fp8_compute and not isinstance(param, FlatParameter))
+                # Unsharded parameters only happens for word_size == 1 or self.fp8_all_gather
+                assert self.world_size == 1 or (self._is_fp8_compute and self.fp8_all_gather and not isinstance(param, FlatParameter))
                 if self.world_size > 1:
                     torch.distributed.all_reduce(
                         to_reduce_grad,
@@ -2290,9 +2292,10 @@ def _prep_grads_for_backward(self) -> None:
         right shape, device, accumulated values, etc.
         """
         for p in self.params:
-            if isinstance(p, FlatParameter) and all(
-                _is_te_module_with_weights(info[1]) for info in p._param_infos
-            ):
+            fused_wgard_accumulation = (self.fp8_all_gather
+                and isinstance(p, FlatParameter)
+                and all(_is_te_module_with_weights(info[1]) for info in p._param_infos))
+            if fused_wgard_accumulation:
                 if getattr(p, "main_grad", None) is None:
                     p.main_grad = torch.empty_like(p, dtype=torch.float32)
                 main_grad_views = p.get_param_views(p.main_grad)
@@ -2367,9 +2370,7 @@ def local_metadata_dict(self) -> Dict[str, Any]:
                         backing_param_name = m.module.flat_param_names[i]
                         names, shapes, numels = m.module.metadata(i)
                     else:
-                        backing_param_name = m._param_name_groups[
-                            m._num_flatten_params
-                        ][i - m._num_flatten_params]
+                        backing_param_name = m._param_name_groups[i][0]
                         names = [backing_param_name]
                         shapes = [p._orig_size]
                         numels = [p._orig_size.numel()]

From 25b23227894d4e5495b58b4783dd737d3e4534db Mon Sep 17 00:00:00 2001
From: Shikai Li <shikaili@fb.com>
Date: Mon, 8 Apr 2024 13:57:14 -0700
Subject: [PATCH 17/20] Fix no pp hanging error.

- Cleans up amax and scale update logic. Amax and scale should be
  done for both weights and parameters. So it should be done at
  forward of each microbatch.

- Consolidate `cast_params` and `all_gather` stream.
---
 .../fully_sharded_data_parallel.py            | 169 +++++++++++-------
 fairscale/nn/misc/flatten_params_wrapper.py   |   3 +-
 2 files changed, 103 insertions(+), 69 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index 3e6923042..49c21097a 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -593,10 +593,9 @@ def should_flatten(name: str) -> bool:
         self.dont_wait_current_stream_for_post_all_gather = False
         self._all_gather_free_event_queue = _FreeEventQueue() if limit_all_gather_events else None
         self._reduce_scatter_free_event_queue = _FreeEventQueue() if limit_reduce_scatter_events else None
-        self.is_full_params_first_rebuilt = True
 
     @property
-    def _is_fp8_compute(self) -> bool:
+    def fp8_compute(self) -> bool:
         return _is_fp8_dtype(self.compute_dtype)
 
     def _get_gradient_predivide_factor(self, world_size: int) -> float:
@@ -832,7 +831,7 @@ def _shard_parameters_(self) -> None:
 
             # If world_size is 1, then we all-reduce grads instead of sharding.
             p._is_sharded = (self.world_size > 1) and (
-                not self._is_fp8_compute or isinstance(p, FlatParameter)
+                not self.fp8_compute or isinstance(p, FlatParameter)
             )
             p._orig_size = p.data.size()
 
@@ -1193,7 +1192,11 @@ def summon_full_params(self, recurse: bool = True, volatile: bool = False) -> Ge
             # Set the state so that we assert when trying to go into
             # forward/backward.
             self.training_state = TrainingState.SUMMON_FULL_PARAMS
-            full_tensors = self._rebuild_full_params(force_full_precision=True)
+            full_tensors = self._rebuild_full_params(
+                force_full_precision=True,
+                wait_for_all_gather=True,
+                is_first_microbatch_fwd=False,
+            )
             assert full_tensors is not None
             with contextlib.ExitStack() as stack:
                 if self.module.is_flattened:
@@ -1286,8 +1289,8 @@ def _param_dtype(self, p: Parameter) -> torch.dtype:
         Returns:
             The dtype to use for the sharded parameters.
         """
-        if self._is_fp8_compute and (not self.fp8_all_gather or
-                                     not isinstance(p, FlatParameter)):
+        if self.fp8_compute and (not self.fp8_all_gather or
+                                 not isinstance(p, FlatParameter)):
             return torch.bfloat16
         else:
             return self.compute_dtype
@@ -1435,8 +1438,6 @@ def _setup_streams(self) -> None:
             return
 
         if torch.cuda.is_available():
-            # Stream to move main FP32 params (may be on CPU) to FP32/FP16/FP8 for forward.
-            self._streams["cast_param"] = torch.cuda.Stream()
             # Stream for all-gathering parameters.
             self._streams["all_gather"] = torch.cuda.Stream()
             # Stream for overlapping grad reduction with the backward pass.
@@ -1471,10 +1472,7 @@ def _wait_for_previous_optim_step(self) -> None:
         """
         if not torch.cuda.is_available():
             return
-        if self.mixed_precision or self.move_params_to_cpu:
-            self._streams["cast_param"].wait_stream(torch.cuda.current_stream())
-        else:
-            self._streams["all_gather"].wait_stream(torch.cuda.current_stream())
+        self._streams["all_gather"].wait_stream(torch.cuda.current_stream())
 
     def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
         self._lazy_init()
@@ -1484,7 +1482,7 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
 
         # For root and mixed precision, we convert the input to FP16 (no_grad is needed for
         # the conversion).
-        is_bf16 = (self.compute_dtype == torch.bfloat16) or self._is_fp8_compute
+        is_bf16 = (self.compute_dtype == torch.bfloat16) or self.fp8_compute
         if self._is_root and self.mixed_precision and self.cast_input:
             args, kwargs = cast_floats_to_right_precision(True, True, is_bf16, *args, **kwargs)
 
@@ -1498,20 +1496,23 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
         if self.force_input_to_fp32 and not self.mixed_precision:
             args, kwargs = cast_floats_to_right_precision(False, False, is_bf16, *args, **kwargs)
 
-        self.module.is_first_batch = not getattr(self, "is_not_first_batch", False)
-        self.is_not_first_batch = True
-
         # All-gather full parameters. This will also transfer FP32 parameters to
         # ``self.compute_dtype`` (e.g., FP16 if *mixed_precision* is ``True``).
-        self._rebuild_full_params()
-        self.is_full_params_first_rebuilt = False
+        self.module.has_unflatten_views = getattr(self.module, "has_unflatten_views", False)
 
+        is_first_microbatch_fwd=kwargs.get("is_first_microbatch", True)
+        self._rebuild_full_params(
+            wait_for_all_gather=True,
+            is_first_microbatch_fwd=is_first_microbatch_fwd
+        )
         if (
             self._fsdp_forward_ordering is not None
-            and self._my_fsdp_instance_idx is not None and self._my_fsdp_instance_idx < len(self._fsdp_forward_ordering) - 1
+            and self._my_fsdp_instance_idx is not None
+            and self._my_fsdp_instance_idx < len(self._fsdp_forward_ordering) - 1
         ):
             self._fsdp_forward_ordering[self._my_fsdp_instance_idx + 1]._rebuild_full_params(
-                wait_for_all_gather=False
+                wait_for_all_gather=False,
+                is_first_microbatch_fwd=is_first_microbatch_fwd
             )
 
         # Register backward hooks to reshard params and reduce-scatter grads.
@@ -1602,13 +1603,18 @@ def _pre_backward_hook(*unused: Any) -> None:
             # idempotent.  So in case they are called unnecessarily, they don't incur much
             # overhead.
             if self.reshard_after_forward:
-                self._rebuild_full_params()
+                self._rebuild_full_params(
+                    wait_for_all_gather=True,
+                    is_first_microbatch_fwd=False,
+                )
                 if (
-                    self.reshard_after_forward
-                    and self._fsdp_forward_ordering is not None
+                    self._fsdp_forward_ordering is not None
                     and self._my_fsdp_instance_idx is not None and self._my_fsdp_instance_idx > 0
                 ):
-                    self._fsdp_forward_ordering[self._my_fsdp_instance_idx - 1]._rebuild_full_params(wait_for_all_gather=False)
+                    self._fsdp_forward_ordering[self._my_fsdp_instance_idx - 1]._rebuild_full_params(
+                        wait_for_all_gather=False,
+                        is_first_microbatch_fwd=False,
+                    )
             else:
                 self._use_full_params()
 
@@ -1784,6 +1790,7 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
         main_grad = getattr(param, "main_grad", None)
         # Only one of `grad` or `main_grad` can exists. Whenever `main_grad is used for accumulation,
         # grad should be set as `None`.
+        assert not param.requires_grad or (grad is not None or main_grad is not None)
         assert not (grad is not None and main_grad is not None)
         # Use `grad` or `main_grad` that is not None and avoid invoking a kernel to check all zeros.
         to_reduce_grad = grad if grad is not None else main_grad
@@ -1800,7 +1807,7 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
             # Free full params.
             self._free_full_params([param])
 
-        if self.mixed_precision and (self._require_backward_grad_sync or self.reshard_after_forward) and not self.fp8_all_gather:
+        if self.mixed_precision and (self._require_backward_grad_sync or self.reshard_after_forward):
             # This is a no-op if reshard_after_forward is True, since we already
             # free the param shard when rebuilding the full params in the
             # pre_backward_hook.
@@ -1873,7 +1880,7 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                 )
             else:
                 # Unsharded parameters only happens for word_size == 1 or self.fp8_all_gather
-                assert self.world_size == 1 or (self._is_fp8_compute and self.fp8_all_gather and not isinstance(param, FlatParameter))
+                assert self.world_size == 1 or (self.fp8_compute and self.fp8_all_gather and not isinstance(param, FlatParameter))
                 if self.world_size > 1:
                     torch.distributed.all_reduce(
                         to_reduce_grad,
@@ -1893,7 +1900,6 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                 release_full_grad_event.record()
                 self._reduce_scatter_free_event_queue.enqueue(release_full_grad_event)
 
-
     def _post_reduction_hook(self, param: Parameter, reduced_grad: torch.Tensor) -> None:
         """Hook to call on each param after the reduce-scatter."""
         assert torch.cuda.current_stream() == self._streams["post_backward"]
@@ -2100,7 +2106,8 @@ def _rebuild_full_params_recursive(self):
             if isinstance(module, FullyShardedDataParallel):
                 module._lazy_init()
                 module._rebuild_full_params(
-                    wait_for_all_gather=False, not_from_recursive=False
+                    wait_for_all_gather=False,
+                    is_first_microbatch_fwd=True,
                 )
 
 
@@ -2109,8 +2116,8 @@ def _rebuild_full_params_recursive(self):
     def _rebuild_full_params(
         self,
         force_full_precision: bool = False,
-        wait_for_all_gather=True,
-        not_from_recursive=True,
+        wait_for_all_gather: bool = True,
+        is_first_microbatch_fwd: bool = False,
     ) -> Optional[List[Tuple[torch.Tensor, bool]]]:
         """
         Gather all shards of params.
@@ -2166,6 +2173,9 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
 
             self.has_full_params = False
 
+        if self.fp8_compute and self.fp8_all_gather:
+            self._update_amax_and_scale_fwd(is_first_microbatch_fwd=is_first_microbatch_fwd)
+
         if self._has_shared_params:
             # self.has_full_params flag can be out of sync if a shared param is
             # sharded by another FSDP instance. An example is that in eval case
@@ -2200,7 +2210,7 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
             if (
                 self.mixed_precision or self.move_params_to_cpu
             ) and not force_full_precision:
-                self._cast_params_for_all_gather(not_from_recursive=not_from_recursive)
+                self._cast_params_for_all_gather()
 
             if self.move_params_to_cpu:
                 if force_full_precision:
@@ -2208,7 +2218,7 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
                     # use pinned memory. Otherwise move p.data to the compute
                     # device.
                     if self.params[0].dtype == self.compute_dtype:
-                        self._cast_params_for_all_gather(not_from_recursive=not_from_recursive)
+                        self._cast_params_for_all_gather()
                     else:
                         for p in self.params:
                             p.data = p.data.to(self.compute_device)
@@ -2322,8 +2332,8 @@ def _free_full_params(self, params: Optional[List[Parameter]] = None) -> None:
         """Free up storage for full parameters."""
         if params is None:
             params = self.params
-        self.is_not_first_batch = False
-        self.is_full_params_first_rebuilt = True
+
+        self.module.has_unflatten_views = False
         self.has_full_params = False
 
         current_stream = torch.cuda.current_stream()
@@ -2478,25 +2488,21 @@ def _use_fp32_param_shard(self, params: Optional[List[Parameter]] = None) -> Non
             p.data = p._fp32_shard
 
     @torch.no_grad()
-    def _cast_params_for_all_gather(
-        self, params: Optional[List[Parameter]] = None,
-        not_from_recursive: bool = False,
-    ) -> None:
-        """Cast FP32 param shard to FP16 for a list of params."""
+    def _update_amax_and_scale_fwd(
+        self,
+        params: Optional[List[Parameter]] = None,
+        is_first_microbatch_fwd: bool = False,
+    ):
+        """Update Amax and scales associated with FP8 parameters."""
         if params is None:
             params = self.params
 
-        with torch.cuda.stream(self._streams["cast_param"]):
+        with torch.cuda.stream(self._streams["all_gather"]):
             for p in params:
-                assert p._fp16_shard is not None
-                alloc_storage_(p._fp16_shard, size=p._fp32_shard.size())
-
-                if self._is_fp8_compute and _is_fp8_dtype(p._fp16_shard.dtype):
+                if self.fp8_compute and _is_fp8_dtype(p._fp16_shard.dtype):
                     assert isinstance(p, FlatParameter)
                     assert len(p._param_infos) == len(p._param_numels)
 
-                    numel_per_shard = p.numel()
-                    offset = -numel_per_shard * self.rank
                     for i in range(len(p._param_infos)):
                         _, m, n = p._param_infos[i]
                         assert _is_te_module_with_weights(m)
@@ -2506,27 +2512,55 @@ def _cast_params_for_all_gather(
                                 num_gemms=2 if isinstance(m, te.LayerNormMLP) else 1
                             )
 
-                        if self.is_full_params_first_rebuilt:
-                            if m.fp8_meta.get("update_amax_and_scale_fwd", False):
-                                if m.fp8_meta["recipe"].reduce_amax:
-                                    FP8GlobalStateManager.copy_amax_from_global_buffer(
-                                        m.fp8_meta, forward=True
-                                    )
-                                    amax_and_scale_update(
-                                        m.fp8_meta,
-                                        True,
-                                        update_weight_scale_inv=True,
-                                    )
-                                    if not_from_recursive:
-                                        FP8GlobalStateManager.set_amax_buffer_key_deletion(
-                                            m.fp8_meta, forward=True
-                                        )
-                                else:
-                                    amax_and_scale_update(
-                                        m.fp8_meta,
-                                        True,
-                                        update_weight_scale_inv=True,
-                                    )
+                        if m.fp8_meta.get("update_amax_and_scale_fwd", False):
+                            if m.fp8_meta["recipe"].reduce_amax:
+                                logging.warning(f"Reduce amax! {is_first_microbatch_fwd}")
+                                FP8GlobalStateManager.copy_amax_from_global_buffer(
+                                    m.fp8_meta, forward=True
+                                )
+                                amax_and_scale_update(
+                                    m.fp8_meta,
+                                    True,
+                                    update_weight_scale_inv=is_first_microbatch_fwd,
+                                )
+                                FP8GlobalStateManager.set_amax_buffer_key_deletion(
+                                    m.fp8_meta, forward=True
+                                )
+                            else:
+                                logging.warning("Not reduce amax! {is_first_microbatch_fwd}")
+                                amax_and_scale_update(
+                                    m.fp8_meta,
+                                    True,
+                                    update_weight_scale_inv=is_first_microbatch_fwd,
+                                )
+                            m.fp8_meta["update_amax_and_scale_fwd"] = False
+
+
+
+    @torch.no_grad()
+    def _cast_params_for_all_gather(
+        self,
+        params: Optional[List[Parameter]] = None,
+    ) -> None:
+        """Cast FP32 params shard to FP16/BF16/FP8 for a list of params."""
+        if params is None:
+            params = self.params
+
+        with torch.cuda.stream(self._streams["all_gather"]):
+            for p in params:
+                assert p._fp16_shard is not None
+                alloc_storage_(p._fp16_shard, size=p._fp32_shard.size())
+
+                if self.fp8_compute and _is_fp8_dtype(p._fp16_shard.dtype):
+                    assert isinstance(p, FlatParameter), "FP8 parameters should be all flatten"
+                    assert len(p._param_infos) == len(p._param_numels)
+
+                    numel_per_shard = p.numel()
+                    offset = -numel_per_shard * self.rank
+                    for i in range(len(p._param_infos)):
+                        _, m, n = p._param_infos[i]
+                        assert _is_te_module_with_weights(m), "Modules with FP8 parameters shoule be TE modules"
+                        assert m.fp8_initialized, "Modules with FP8 parameters should be initialized with scales"
 
                         numel = p._param_numels[i]
                         if offset + numel <= 0 or offset >= numel_per_shard:
@@ -2563,7 +2597,6 @@ def _cast_params_for_all_gather(
                     )
                     p.data = p._fp16_shard
 
-        self._streams["all_gather"].wait_stream(self._streams["cast_param"])
 
     @torch.no_grad()
     def _free_fp16_param_shard(self, params: Optional[List[Parameter]] = None) -> None:
diff --git a/fairscale/nn/misc/flatten_params_wrapper.py b/fairscale/nn/misc/flatten_params_wrapper.py
index ae2f3d792..da947dc31 100644
--- a/fairscale/nn/misc/flatten_params_wrapper.py
+++ b/fairscale/nn/misc/flatten_params_wrapper.py
@@ -486,8 +486,9 @@ def load_state_dict(
             return super().load_state_dict(state_dict, strict)
 
     def forward(self, *inputs: Any, **kwinputs: Any) -> Any:
-        if getattr(self, "is_first_batch", False):
+        if not getattr(self, "has_unflatten_views", False):
             self._unflatten_params_as_views()
+            self.has_unflatten_views = True
         return self.module(*inputs, **kwinputs)
 
     def get_param_views(self, external_data_list: Optional[List[Optional[Tensor]]] = None) -> Iterator[Tensor]:

From 0d1502b8979fd64dfbf05aeae6ca86e9534b340e Mon Sep 17 00:00:00 2001
From: Shikai Li <shikaili@fb.com>
Date: Wed, 10 Apr 2024 13:16:41 -0700
Subject: [PATCH 18/20] Clean up shard offset calculation logic.

---
 .../fully_sharded_data_parallel.py            | 30 ++++++++++++-------
 1 file changed, 19 insertions(+), 11 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index 49c21097a..e7f12eba9 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -444,7 +444,7 @@ def __init__(
         self.flatten_parameters = flatten_parameters
         self.move_params_to_cpu = move_params_to_cpu or cpu_offload
         self.compute_dtype = compute_dtype or (torch.float16 if mixed_precision else torch.float32)
-        self.buffer_dtype = buffer_dtype or self.compute_dtype
+        self.buffer_dtype = buffer_dtype or (torch.bfloat16 if _is_fp8_dtype(self.compute_dtype) else self.compute_dtype)
         self.move_grads_to_cpu = self.move_params_to_cpu if move_grads_to_cpu is None else move_grads_to_cpu
         self.bucket_cap_mb = bucket_cap_mb
         self.compute_device = compute_device or _get_default_cuda_device(module)
@@ -2514,7 +2514,6 @@ def _update_amax_and_scale_fwd(
 
                         if m.fp8_meta.get("update_amax_and_scale_fwd", False):
                             if m.fp8_meta["recipe"].reduce_amax:
-                                logging.warning(f"Reduce amax! {is_first_microbatch_fwd}")
                                 FP8GlobalStateManager.copy_amax_from_global_buffer(
                                     m.fp8_meta, forward=True
                                 )
@@ -2527,7 +2526,6 @@ def _update_amax_and_scale_fwd(
                                     m.fp8_meta, forward=True
                                 )
                             else:
-                                logging.warning("Not reduce amax! {is_first_microbatch_fwd}")
                                 amax_and_scale_update(
                                     m.fp8_meta,
                                     True,
@@ -2552,28 +2550,38 @@ def _cast_params_for_all_gather(
                 alloc_storage_(p._fp16_shard, size=p._fp32_shard.size())
 
                 if self.fp8_compute and _is_fp8_dtype(p._fp16_shard.dtype):
+                    assert p._is_sharded
                     assert isinstance(p, FlatParameter), "FP8 parameters should be all flatten"
                     assert len(p._param_infos) == len(p._param_numels)
 
                     numel_per_shard = p.numel()
-                    offset = -numel_per_shard * self.rank
+
+                    flat_index = 0
+                    flat_begin = numel_per_shard * self.rank
+                    flat_end = flat_begin + numel_per_shard
+
                     for i in range(len(p._param_infos)):
                         _, m, n = p._param_infos[i]
+
                         assert _is_te_module_with_weights(m), "Modules with FP8 parameters shoule be TE modules"
                         assert m.fp8_initialized, "Modules with FP8 parameters should be initialized with scales"
 
                         numel = p._param_numels[i]
-                        if offset + numel <= 0 or offset >= numel_per_shard:
-                            offset += numel
+
+                        if flat_index >= flat_end:
+                            break
+                        shard_begin = max(flat_index - flat_begin, 0)
+
+                        flat_index += numel
+                        if flat_index <= flat_begin:
                             continue
+                        shard_end = min(flat_index - flat_begin, numel_per_shard)
 
                         fp8_dtype_forward = te.fp8.get_fp8_te_dtype(
                             m.fp8_meta["recipe"], fprop_tensor=True
                         )
-                        begin = max(offset, 0)
-                        end = min(offset + numel, numel_per_shard)
                         cast_to_fp8(
-                            p._fp32_shard[begin:end].bfloat16(),
+                            p._fp32_shard[shard_begin:shard_end].bfloat16().contiguous(),
                             m.fp8_meta["scaling_fwd"],
                             (
                                 FP8FwdTensors.GEMM2_WEIGHT
@@ -2581,9 +2589,9 @@ def _cast_params_for_all_gather(
                                 else FP8FwdTensors.GEMM1_WEIGHT
                             ),
                             fp8_dtype_forward,
-                            out=p._fp16_shard[begin:end],
+                            out=p._fp16_shard[shard_begin:shard_end],
                         )
-                        offset += numel
+                    # Doesn't need to set padding elements.
                     p.data = p._fp16_shard.view(
                             torch.float8_e4m3fn
                             if fp8_dtype_forward == DType.kFloat8E4M3

From 5edb109889ea090bc20eeaa015c54a120434eca6 Mon Sep 17 00:00:00 2001
From: Shikai Li <shikaili@fb.com>
Date: Mon, 15 Apr 2024 23:20:47 -0700
Subject: [PATCH 19/20] Unify compute dtype setting.

---
 .../fully_sharded_data_parallel.py            | 82 ++++++++++---------
 1 file changed, 44 insertions(+), 38 deletions(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index e7f12eba9..71e98d8a3 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -443,8 +443,8 @@ def __init__(
         self.fp8_all_gather = fp8_all_gather
         self.flatten_parameters = flatten_parameters
         self.move_params_to_cpu = move_params_to_cpu or cpu_offload
-        self.compute_dtype = compute_dtype or (torch.float16 if mixed_precision else torch.float32)
-        self.buffer_dtype = buffer_dtype or (torch.bfloat16 if _is_fp8_dtype(self.compute_dtype) else self.compute_dtype)
+        self.compute_dtype = compute_dtype or (torch.bfloat16 if mixed_precision else torch.float32)
+        self.buffer_dtype = buffer_dtype or self.compute_dtype
         self.move_grads_to_cpu = self.move_params_to_cpu if move_grads_to_cpu is None else move_grads_to_cpu
         self.bucket_cap_mb = bucket_cap_mb
         self.compute_device = compute_device or _get_default_cuda_device(module)
@@ -489,6 +489,11 @@ def __init__(
                 param_names.append(param_name)
                 params.append(param)
 
+        for m in module.modules():
+            for param in m.parameters():
+                if not getattr(param, "_is_te_param", False):
+                    param._is_te_param = _is_te_module_with_weights(m)
+
         self._has_params = len(params) > 0
         self._has_shared_params = False
 
@@ -505,27 +510,34 @@ def __init__(
 
         # For now, it is either all flatten or none flatten. This will be extended to
         # multiple flatten groups in my next PR.
-        def should_flatten(name: str) -> bool:
+        no_te_params = not any(p._is_te_param for p in params)
+        def should_flatten(name: str, param: Parameter) -> bool:
+            if not self.flatten_parameters:
+                return False
+            # If no TE weights or no FP8 AllGather, then flatten them all. Shard as compute dtype.
+            if no_te_params or not self.fp8_all_gather:
+                return True
             # `*_norm_weights` are numerics-sensitive and cannot be quantized to fp8.
-            return self.flatten_parameters and (not self.fp8_all_gather or "norm_weight" not in name)
+            return  param._is_te_param and "norm_weight" not in name
+
+        to_be_flatten_param_names = []
+        to_be_flatten_params = []
+        non_flatten_param_names = []
+        non_flatten_params = []
+        for name, param in zip(param_names, params):
+            if should_flatten(name, param):
+                to_be_flatten_param_names.append(name)
+                to_be_flatten_params.append(param)
+            else:
+                non_flatten_param_names.append(name)
+                non_flatten_params.append(param)
 
-        to_be_flatten_params: List[List[Parameter]] = [
-            param
-            for param, name in zip(params, param_names)
-            if should_flatten(name)
-        ]
         if to_be_flatten_params:
             to_be_flatten_params = [to_be_flatten_params]
-        non_flatten_params: List[List[Parameter]] = [
-            param
-            for param, name in zip(params, param_names)
-            if not should_flatten(name)
-        ]
-        param_name_groups: List[List[str]] = [
-            [n for n in param_names if should_flatten(n)]
-        ] + [
-            [n] for n in param_names if not should_flatten(n)
-        ]
+
+        param_name_groups: List[List[str]] = [to_be_flatten_param_names] + [[n] for n in non_flatten_param_names]
+
+        logging.info(f"param_names: {param_name_groups}")
         del param_names
 
         self._fsdp_wrapped_module: nn.Module = FlattenParamsWrapper(
@@ -533,6 +545,9 @@ def should_flatten(name: str) -> bool:
         )
         del module  # free original module in case it helps garbage collection
 
+        for param in self._fsdp_wrapped_module.flat_params:
+            param._is_fp8_param = not no_te_params and self.fp8_all_gather
+
         # Now, in this FSDP wrapper class, we keep a list of to-be-flatten and not-to-be-flatten
         # params for doing sharding, gradient hooks, etc. Note, the ordering of the
         # list matters: flatten params are always in the front.
@@ -594,10 +609,6 @@ def should_flatten(name: str) -> bool:
         self._all_gather_free_event_queue = _FreeEventQueue() if limit_all_gather_events else None
         self._reduce_scatter_free_event_queue = _FreeEventQueue() if limit_reduce_scatter_events else None
 
-    @property
-    def fp8_compute(self) -> bool:
-        return _is_fp8_dtype(self.compute_dtype)
-
     def _get_gradient_predivide_factor(self, world_size: int) -> float:
         factor: int = 1
         while world_size % factor == 0 and world_size / factor > factor:
@@ -830,9 +841,7 @@ def _shard_parameters_(self) -> None:
                 assert p.dtype == torch.float32
 
             # If world_size is 1, then we all-reduce grads instead of sharding.
-            p._is_sharded = (self.world_size > 1) and (
-                not self.fp8_compute or isinstance(p, FlatParameter)
-            )
+            p._is_sharded = (self.world_size > 1) and isinstance(p, FlatParameter)
             p._orig_size = p.data.size()
 
             if not p._is_sharded:
@@ -1289,11 +1298,9 @@ def _param_dtype(self, p: Parameter) -> torch.dtype:
         Returns:
             The dtype to use for the sharded parameters.
         """
-        if self.fp8_compute and (not self.fp8_all_gather or
-                                 not isinstance(p, FlatParameter)):
-            return torch.bfloat16
-        else:
-            return self.compute_dtype
+        if getattr(p, "_is_fp8_param", False):
+            return torch.float8_e4m3fn
+        return self.compute_dtype
 
     @torch.no_grad()
     def _init_param_attributes(self, p: Parameter) -> None:
@@ -1482,7 +1489,7 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
 
         # For root and mixed precision, we convert the input to FP16 (no_grad is needed for
         # the conversion).
-        is_bf16 = (self.compute_dtype == torch.bfloat16) or self.fp8_compute
+        is_bf16 = self.compute_dtype == torch.bfloat16
         if self._is_root and self.mixed_precision and self.cast_input:
             args, kwargs = cast_floats_to_right_precision(True, True, is_bf16, *args, **kwargs)
 
@@ -1514,7 +1521,6 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
                 wait_for_all_gather=False,
                 is_first_microbatch_fwd=is_first_microbatch_fwd
             )
-
         # Register backward hooks to reshard params and reduce-scatter grads.
         # These need to be re-registered every forward pass.
         self._register_post_backward_hooks()
@@ -1880,7 +1886,7 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
                 )
             else:
                 # Unsharded parameters only happens for word_size == 1 or self.fp8_all_gather
-                assert self.world_size == 1 or (self.fp8_compute and self.fp8_all_gather and not isinstance(param, FlatParameter))
+                assert self.world_size == 1 or not isinstance(param, FlatParameter)
                 if self.world_size > 1:
                     torch.distributed.all_reduce(
                         to_reduce_grad,
@@ -2173,7 +2179,7 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
 
             self.has_full_params = False
 
-        if self.fp8_compute and self.fp8_all_gather:
+        if self.fp8_all_gather:
             self._update_amax_and_scale_fwd(is_first_microbatch_fwd=is_first_microbatch_fwd)
 
         if self._has_shared_params:
@@ -2499,7 +2505,7 @@ def _update_amax_and_scale_fwd(
 
         with torch.cuda.stream(self._streams["all_gather"]):
             for p in params:
-                if self.fp8_compute and _is_fp8_dtype(p._fp16_shard.dtype):
+                if _is_fp8_dtype(p._fp16_shard.dtype):
                     assert isinstance(p, FlatParameter)
                     assert len(p._param_infos) == len(p._param_numels)
 
@@ -2549,7 +2555,7 @@ def _cast_params_for_all_gather(
                 assert p._fp16_shard is not None
                 alloc_storage_(p._fp16_shard, size=p._fp32_shard.size())
 
-                if self.fp8_compute and _is_fp8_dtype(p._fp16_shard.dtype):
+                if _is_fp8_dtype(p._fp16_shard.dtype):
                     assert p._is_sharded
                     assert isinstance(p, FlatParameter), "FP8 parameters should be all flatten"
                     assert len(p._param_infos) == len(p._param_numels)
@@ -2563,7 +2569,7 @@ def _cast_params_for_all_gather(
                     for i in range(len(p._param_infos)):
                         _, m, n = p._param_infos[i]
 
-                        assert _is_te_module_with_weights(m), "Modules with FP8 parameters shoule be TE modules"
+                        assert _is_te_module_with_weights(m), f"Modules {m} with FP8 parameters shoule be TE modules"
                         assert m.fp8_initialized, "Modules with FP8 parameters should be initialized with scales"
 
                         numel = p._param_numels[i]

From 2df199f89d3eea66cf40bf4732111b1c3972e470 Mon Sep 17 00:00:00 2001
From: Shikai Li <shikaili@fb.com>
Date: Wed, 17 Apr 2024 16:45:15 -0700
Subject: [PATCH 20/20] Have FP16 and FP8 sharded in the same way.

---
 fairscale/nn/data_parallel/fully_sharded_data_parallel.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
index 71e98d8a3..90ae056df 100644
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -515,7 +515,7 @@ def should_flatten(name: str, param: Parameter) -> bool:
             if not self.flatten_parameters:
                 return False
             # If no TE weights or no FP8 AllGather, then flatten them all. Shard as compute dtype.
-            if no_te_params or not self.fp8_all_gather:
+            if no_te_params: #or not self.fp8_all_gather:
                 return True
             # `*_norm_weights` are numerics-sensitive and cannot be quantized to fp8.
             return  param._is_te_param and "norm_weight" not in name