Revert "[Megatron-FSDP] Add dtype customization to Megatron-FSDP. (#3067)"

This reverts commit b969f76.

Signed-off-by: oliver könig <okoenig@nvidia.com>

Author: ko3n1g

megatron/core/distributed/distributed_data_parallel_config.py

@@ -33,10 +33,7 @@ class DistributedDataParallelConfig:
     """
 
     check_for_nan_in_grad: bool = False
-    """
-    If true, check for NaNs and Infs in gradients _before_ communication collective.
-    Invoked by `start_grad_sync` such as in the Megatron-LM DDP training API.
-    """
+    """If true, check for NaNs and Infs in gradients _before_ communication collective."""
 
     check_for_large_grads: bool = False
     """If true, check for unexpectedly large gradients _before_ communication collective."""
@@ -81,7 +78,7 @@ class DistributedDataParallelConfig:
 
     data_parallel_sharding_strategy: str = 'no_shard'
     """Sharding strategy for FSDP. Valid values are 'no_shard', 'optim',
-      'optim_grads', 'optim_grads_params'."""
+        'optim_grads', 'optim_grads_params'."""
 
     gradient_reduce_div_fusion: bool = True
     """If true, perform gradient reduce and division fusion."""
@@ -94,6 +91,9 @@ class DistributedDataParallelConfig:
       disables prefetching and may degrade performance. Adjust this value
       based on your system's memory and performance requirements."""
 
+    preserve_fp32_weights: bool = True
+    """If true, preserve fp32 weights in the Megatron FSDP ParamAndGradBuffer."""
+
     keep_fp8_transpose_cache: bool = False
     """If true, keep the fp8 transpose cache when using Megatron FSDP."""
 
@@ -128,7 +128,7 @@ class DistributedDataParallelConfig:
        allocated buffer for the bucket that does not fit, it will enable NCCL 
        user buffer with the cost of more memory usage. If false, FSDP will use
        Dynamic memory allocator, NCCL user buffer won't not enabled, which 
-       usually leads to low performance.
+       usually leads to low performance. 
     """
 
     fsdp_all_gather_in_start_param_sync: bool = True
@@ -142,7 +142,8 @@ class DistributedDataParallelConfig:
     outer_dp_sharding_strategy: str = 'no_shard'
     """
     Sharding strategy for outer data parallel group in Hybrid Sharded Data Parallel (HSDP) mode.
-    Valid values are 'no_shard', 'optim'. This option is only effective when Hybrid FSDP is enabled.
+    Valid values are 'no_shard', 'optim', 'optim_grads', 'optim_grads_params'.
+    This option is only effective when Hybrid FSDP is enabled.
     """
 
     disable_symmetric_registration: bool = False

megatron/core/distributed/fsdp/mcore_fsdp_adapter.py

@@ -45,11 +45,7 @@
 from megatron.core.utils import is_te_min_version, log_single_rank
 
 try:
-    from megatron.core.distributed.fsdp.src.megatron_fsdp import (
-        FSDPDistributedIndex,
-        MegatronFSDP,
-        MixedPrecisionPolicy,
-    )
+    from megatron.core.distributed.fsdp.src.megatron_fsdp import FSDPDistributedIndex, MegatronFSDP
 
     HAVE_MEGATRON_FSDP = True
 except ImportError as import_megatron_fsdp_error:
@@ -70,9 +66,6 @@ def __init__(
         ddp_config: DistributedDataParallelConfig,
         module: torch.nn.Module,
         fsdp_unit_modules: Optional[List[torch.nn.Module]] = None,
-        main_params_dtype: Optional[torch.dtype] = torch.float32,
-        main_grads_dtype: Optional[torch.dtype] = torch.float32,
-        grad_comm_dtype: Optional[torch.dtype] = torch.float32,
         disable_bucketing: bool = False,
         device: Optional[torch.device] = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
@@ -89,17 +82,6 @@ def __init__(
             logging.INFO,
             f'Setting up DistributedDataParallel with config {self.ddp_config}',
         )
-        self.mp_policy = MixedPrecisionPolicy(
-            main_params_dtype=main_params_dtype,
-            # Grandfathered Argument: grad_reduce_in_fp32
-            main_grads_dtype=torch.float32 if ddp_config.grad_reduce_in_fp32 else main_grads_dtype,
-            grad_comm_dtype=grad_comm_dtype,
-        )
-        log_single_rank(
-            logger,
-            logging.INFO,
-            f'Setting up Megatron-FSDP MixedPrecisionPolicy with config {self.mp_policy}',
-        )
 
         self.megatron_fsdp_dist_index = self._init_dist_index(pg_collection)
 
@@ -128,7 +110,6 @@ def __init__(
             config=config,
             module=MegatronFSDP(
                 ddp_config=ddp_config,
-                mixed_precision_policy=self.mp_policy,
                 module=module,
                 fsdp_unit_modules=self.fsdp_unit_modules,
                 disable_bucketing=disable_bucketing,

megatron/core/distributed/fsdp/src/README.md

@@ -129,7 +129,6 @@ import torch
 from megatron_fsdp import (
     fully_shard_model,
     fully_shard_optimizer,
-    MixedPrecisionPolicy,
 )
 ```
 
@@ -197,8 +196,10 @@ model = fully_shard_model(
     outer_dp_sharding_strategy=1,
     # Initialize the model on devices in shards to avoid OOM. Requires device("meta")-init for model.
     init_model_with_meta_device=True,
-    # Mixed-Precision Policy for controlling compute and communication precision in Megatron-FSDP.
-    mixed_precision_policy=MixedPrecisionPolicy(),
+    # Reduce gradients in FP32.
+    grad_reduce_in_fp32=False,
+    # Store distributed optimization state in FP32.
+    preserve_fp32_weights=True,
     # Sync parameters and gradients each step. Allows for gradient transformations after backward pass,
     # and synchronizes parameters and gradients across HSDP groups, but deactivates compute-communication
     # overlap going into the subsequent training step.
@@ -284,7 +285,7 @@ Megatron-FSDP's `fully_shard_*` API has a comprehensive set of arguments for fin
 
 - `fsdp_unit_modules` is a list of sub-module classes or `str` import-paths associated with modules that you want `MegatronFSDP` to fully-shard.
   - Required if `1`, `2`, or `3` are specified as the sharding strategy. Defaults to `None`, in which case Megatron-FSDP will replicate the parameters similar to DDP.
-- `zero_dp_strategy` (and `outer_dp_sharding_strategy`) configure different degrees of zero-redundancy data parallelism as described in [ZeRO (Zero Redundancy Optimizer)](https://arxiv.org/abs/1910.02054). It reduces CUDA memory utilization during model training by distributing model parameters, gradients, and optimizer states across multiple devices in the DP `ProcessGroup`, and collectively communicating subsets of parameters and gradients to specific devices when needed for computation or differentiation. More aggressive sharding strategies will entail more communication overhead, with `no_shard` being the least memory efficient but most communication efficient, and `optim_grads_params` being the most memory efficient but least communication efficient. Additionally, `outer_dp_sharding_strategy` supports `no_shard` ([Hybrid-Sharded Data Parallelism (HSDP)](https://arxiv.org/pdf/2304.11277)) and `optim` (`HFSDP` = Fully-Sharded Optimizer State + _HSDP_, requires `zero_dp_strategy='optim_grads_params'`), after specifying the "outer" DP group (`dp_outer_dim` / `hybrid_fsdp_group`).
+- `zero_dp_strategy` (and `outer_dp_sharding_strategy`) configure different degrees of zero-redundancy data parallelism as described in [ZeRO (Zero Redundancy Optimizer)](https://arxiv.org/abs/1910.02054). It reduces CUDA memory utilization during model training by distributing model parameters, gradients, and optimizer states across multiple devices in the DP `ProcessGroup`, and collectively communicating subsets of parameters and gradients to specific devices when needed for computation or differentiation. More aggressive sharding strategies will entail more communication overhead, with `no_shard` being the least memory efficient but most communication efficient, and `optim_grads_params` being the most memory efficient but least communication efficient. `outer_dp_sharding_strategy` has the same options, except for the (required) "outer" DP group (`dp_outer_dim` / `hybrid_fsdp_group`) when using [Hybrid-Sharded Data Parallelism (HSDP)](https://arxiv.org/pdf/2304.11277), and only `no_shard` (DP Replication) and `optim` (Optimizer State Hybrid Sharding, requires `zero_dp_strategy='optim_grads_params`) are supported.
   - Default: `optim_grads_params` or `3` for `zero_dp_strategy` and `no_shard` or `0` for `outer_dp_sharding_strategy`
   - `0` or `no_shard` implies that your model is not sharded. Similar memory usage to `DDP`.
   - `1` or `optim` implies that your optimizer state is sharded for distributed optimization. Similar to optimizer state sharding in `ZeRO-DP`.
@@ -304,25 +305,16 @@ Megatron-FSDP's `fully_shard_*` API has a comprehensive set of arguments for fin
 - `init_model_with_meta_device` has `MegatronFSDP` initialize your `meta`-device model in shards on every CUDA device to avoid OOM when initializing extremely large models that cannot fit on a single device. Users can initialize their model on a [`meta`-device](https://docs.pytorch.org/docs/stable/meta.html) (`with torch.device('meta'): ...`), and ``MegatronFSDP`` will further shard and initialize the model parameters layer-by-layer adhering to the customizable `module.reset_parameters` method, which prevents the entire model from being allocated in memory at any point during runtime.
     - Defaults to `False`.
     - Note that the `device` argument which installs your model on a specific device or rank will be deactivated when `init_model_with_meta_device=True`.
-- `mixed_precision_policy` takes a `megatron_fsdp.MixedPrecisionPolicy` that configures mixed-precision compute and communication for Megatron-FSDP. Configuration options include:
-    - `main_params_dtype` controls the data-type for parameters used in distributed optimization or quantization. 
-        - Defaults to `torch.float32`.
-        - If set to `None`, the native model compute parameter data-type will be utilized.
-        - Requires specification (cannot be `None`) when using `FP8` parameters with Megatron-FSDP.
-    - `main_grads_dtype` controls the data-type for gradients used in distributed optimization.
-        - Defaults to `torch.float32`, which is highly-recommended for accuracy at scale.
-        - If set to `None`, the model native gradient data-type will be utilized.
-    - `grad_comm_dtype` controls the data-type for gradient communications (RS / AR) when reducing gradients. Lower precision `grad_comm_dtype` improves (communication) performance, but may increase memory utilization or sacrifice gradient precision in certain cases.
-        - Defaults to `torch.float32`.
-        - If set to `None`, the `main_grads_dtype` data-type will be utilized.
-        - If using `no_shard`, `optim`, or a `FixedPoolAllocator` (`fsdp_double_buffer`), allocating `dtype`-custom gradient communication buffers (per FSDP group) adds memory overhead.
-        - If using NCCL UBR v2.27+ (`nccl_ub=True`), gradient reduction may be performed in high-precision depending on the network domain (NVLink or IB), and can enable mixed-precision communication and accumulation, e.g. setting grad_comm_dtype to BF16 can support FP32 reduction even though we have BF16 input and output communication buffers. Otherwise, gradients will be reduced in `grad_comm_dtype` (and accumulated in `main_grads_dtype`) as usual.
+- `grad_reduce_in_fp32` will reduce gradients in `FP32` precision (in contrast to the lower `BF16` or `FP8` model training precision).
+    - Defaults to `False`.
+    - `torch.distributed.fsdp.MixedPrecisionPolicy` will be supported in the near future.
+- `preserve_fp32_weights` will preserve a `FP32` precision version of model parameters utilized for optimization.
+    - Defaults to `True`.
+    - `torch.distributed.fsdp.MixedPrecisionPolicy` will be supported in the near future.
 - `overlap_grad_reduce` and `overlap_param_gather` will overlap gradient [`reduce-scatter`](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html#reducescatter) and parameter [`all-gather`](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html#allgather) group communications with backward and forward compute with asynchronous calls and pre-fetching. (In the case of `no_shard`, parameters are not gathered but gradient [`all-reduce`](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html#allreduce) is overlapped.)
     - Both default to `True`.
-- `sync_model_each_microbatch` will trigger a `wait` (`MegatronFSDP.finish_grad_sync()`) on gradient reduction, parameter de-allocation, and optimizer parameter / gradient installation (in preparation for `optimizer.step()`) after every forward-backward pass. When using HSDP, parameters and gradients will be all-gathered and reduced respectively on the "outer" DP group each training step instead of each optimization cycle. This behavior is desirable for a transparent and user-friendly sharded training loop where post-backward transformations on the gradient and a clean compute / memory state are necessary within and between training iterations, but damages performance in situations where optimization is delayed (e.g. gradient accumulation) when the communications of the previous training iteration can be overlapped with the compute of the next training iteration. Will also override `is_last_microbatch` / `microbatch_count` logic in `MegatronFSDP`.
+- `sync_model_each_microbatch` will trigger a `wait` (`MegatronFSDP.finish_grad_sync()`) on gradient reduction, parameter de-allocation, and optimizer parameter / gradient installation (in preparation for `optimizer.step()`) after every forward-backward pass. When using HSDP, parameters and gradients will be all-gathered and reduced respectively on the "outer" DP group each training step instead of each optimization cycle. This behavior is desirable for a transparent and user-friendly sharded training loop where post-backward transformations on the gradient and a clean compute / memory state are necessary between training iterations, but damages performance in situations where optimization is delayed (e.g. gradient accumulation) where the communications of the previous training iteration can be overlapped with the compute of the next training iteration. Will also override `is_last_microbatch` / `microbatch_count` logic in `MegatronFSDP`.
     - Defaults to `True` for `fully_shard`, but defaults to `False` when using the `MegatronFSDP` class directly.
-    - Can also be controlled with the `MegatronFSDP.sync()` context manager, or through invoking `MegatronFSDP.set_model_auto_sync(bool)`.
-    - WARNING: When this synchronization feature is activated in conjunction with `no_shard` / `0` or `optim` / `1` sharding strategies, the user is responsible for calling `MegatronFSDP.zero_grad_buffer()` or `optimizer.zero_grad()` after the subsequent forward-backward pass. This is because un-sharded gradients are all-reduced directly into the gradient accumulation buffer, and this buffer should not be all-reduced more than once per optimization cycle! Analogous to the justification for the [`no_sync()` API for PyTorch DistributedDataParallel](https://docs.pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel.no_sync).
 - `enable_fine_grained_param_gather` modifies FSDP to all-gather parameters with per-Module granularity instead of collectively unsharding all sub-modules of a unit module in Megatron-FSDP.
     - Defaults to `False`.
 - `keep_fp8_transpose_cache` will keep the fp8 transpose cache when using `MegatronFSDP`. This option will cause (number of parameter $\times$ 1 Byte) of memory overhead, but can skip the weight transpose operation in the backward propagation. This feature will not give any benefit from the Blackwell architecture.
@@ -370,10 +362,8 @@ mfsdp_model = fully_shard_model(
     fsdp_unit_modules=[te.pytorch.TransformerLayer],
     # Only FSDP / ZeRO-3 supports FP8 parameters.
     zero_dp_strategy=3,
-    # FP32 main weights needed for FP8 parameters.
-    mixed_precision_policy=MixedPrecisionPolicy(
-        main_params_dtype=torch.float32
-    ),
+    # Needed for FP8 parameters. (Default is already True.)
+    preserve_fp32_weights=True,
     # Needed for select FP8 recipes.
     keep_fp8_transpose_cache=True,
 )

megatron/core/distributed/fsdp/src/megatron_fsdp/__init__.py

@@ -15,7 +15,6 @@
 from .distributed_data_parallel_config import DistributedDataParallelConfig
 from .fully_shard import fully_shard, fully_shard_model, fully_shard_optimizer
 from .megatron_fsdp import MegatronFSDP
-from .mixed_precision import MixedPrecisionPolicy
 from .package_info import (
     __contact_emails__,
     __contact_names__,
@@ -35,7 +34,6 @@
     "DistributedDataParallelConfig",
     "MegatronFSDP",
     "FSDPDistributedIndex",
-    "MixedPrecisionPolicy",
     "fully_shard",
     "fully_shard_model",
     "fully_shard_optimizer",