Clean and fin

Author: muellerzr

benchmarks/fp8/ms_amp/Dockerfile

@@ -1,11 +1,11 @@
 FROM ghcr.io/azure/msamp
 
 RUN pip install transformers evaluate datasets
-# RUN git clone https://github.com/huggingface/accelerate
+RUN git clone https://github.com/huggingface/accelerate
 
-# RUN cd accelerate && \
-#     pip install -e . && \
-#     cd benchmarks/fp8
+RUN cd accelerate && \
+    pip install -e . && \
+    cd benchmarks/fp8
 
 CMD ["bash"]
 

benchmarks/fp8/ms_amp/distrib_deepspeed.py

@@ -16,32 +16,27 @@
 This script tests to ensure that `accelerate` performs at the same level as raw `MS-AMP`.
 
 This particular script verifies this for DeepSpeed training.
+
+NOTE: MS-AMP does *not* support ZeRO-3.
 """
-from unittest.mock import patch
 
-from msamp import deepspeed
+# import msamp.deepspeed as msamp_deepspeed
 import evaluate
 import torch
-# import transformer_engine.common.recipe as te_recipe
-# import transformer_engine.pytorch as te
 from fp8_utils import evaluate_model, get_training_utilities
-# from transformer_engine.common.recipe import DelayedScaling
+from msamp import deepspeed as msamp_deepspeed
 
 from accelerate import Accelerator, DeepSpeedPlugin
 from accelerate.state import AcceleratorState
-from accelerate.utils import FP8RecipeKwargs, set_seed
+from accelerate.utils import set_seed
 
 
 MODEL_NAME = "bert-base-cased"
 METRIC = evaluate.load("glue", "mrpc")
 
 
 def train_baseline(zero_stage: int = 1, opt_level: str = "O1"):
-    # This forces transformers to think Zero-3 Init should be used
-    with patch("transformers.integrations.deepspeed.is_deepspeed_zero3_enabled") as mock:
-        mock.return_value = zero_stage == 3
     set_seed(42)
-
     accelerator = Accelerator()
     model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(
         MODEL_NAME, accelerator=accelerator
@@ -57,7 +52,6 @@ def train_baseline(zero_stage: int = 1, opt_level: str = "O1"):
             "stage": zero_stage,
             "offload_optimizer": {"device": "none", "nvme_path": None},
             "offload_param": {"device": "none", "nvme_path": None},
-            "stage3_gather_16bit_weights_on_model_save": False,
         },
         "gradient_clipping": 1.0,
         "steps_per_print": np.inf,
@@ -67,15 +61,14 @@ def train_baseline(zero_stage: int = 1, opt_level: str = "O1"):
         "msamp": {
             "enabled": True,
             "opt_level": opt_level,
-        }
+        },
     }
-
     (
         model,
         optimizer,
         _,
         _,
-    ) = deepspeed.initialize(
+    ) = msamp_deepspeed.initialize(
         model=model,
         optimizer=optimizer,
         config_params=config,
@@ -107,18 +100,14 @@ def train_baseline(zero_stage: int = 1, opt_level: str = "O1"):
     return base_model_results, trained_model_results
 
 
-def train_integration(zero_stage: int = 1):
+def train_integration(zero_stage: int = 1, opt_level: str = "O1"):
     set_seed(42)
-    FP8_RECIPE_KWARGS = {"fp8_format": "HYBRID", "amax_history_len": 32, "amax_compute_algo": "max"}
-    kwargs_handlers = [FP8RecipeKwargs(backend="TE", **FP8_RECIPE_KWARGS)]
-    AcceleratorState()._reset_state(True)
     deepspeed_plugin = DeepSpeedPlugin(
         zero_stage=zero_stage,
-        zero3_init_flag=zero_stage == 3,
-    )
-    accelerator = Accelerator(
-        mixed_precision="fp8", kwargs_handlers=kwargs_handlers, deepspeed_plugin=deepspeed_plugin
+        enable_msamp=True,
+        msamp_opt_level=opt_level,
     )
+    accelerator = Accelerator(mixed_precision="fp8", deepspeed_plugin=deepspeed_plugin)
     accelerator.state.deepspeed_plugin.deepspeed_config["train_micro_batch_size_per_gpu"] = 16
 
     model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(
@@ -128,13 +117,9 @@ def train_integration(zero_stage: int = 1):
     model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
     base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
     model.train()
-    model_outputs = []
-    data = []
     for _ in range(2):
         for batch in train_dataloader:
             outputs = model(**batch)
-            data.append(batch.to("cpu"))
-            model_outputs.append(outputs.logits.to("cpu"))
             loss = outputs.loss
             accelerator.backward(loss)
             optimizer.step()
@@ -151,32 +136,26 @@ def train_integration(zero_stage: int = 1):
         trained_model_results["f1"] > base_model_results["f1"]
     ), f'F1 score should be higher for the trained model: {trained_model_results["f1"]} > {base_model_results["f1"]}'
 
+    AcceleratorState()._reset_state(True)
     return base_model_results, trained_model_results
 
 
 if __name__ == "__main__":
-    # results = {"1": [], "2": [], "3": []}
-    # for zero_stage in [1, 2, 3]:
-    #     for opt_level in ["O1", "O2", "O3"]:
-    baseline_not_trained, baseline_trained = train_baseline(3, "O3")
-    print(baseline_not_trained, baseline_trained)
-    # results[str(zero_stage)].append({"opt_level": opt_level, "not_trained": baseline_not_trained, "trained": baseline_trained})
-    # for stage, stage_results in results.items():
-    #     print(f'zero_stage={stage}:\n')
-    #     for result in stage_results:
-    #         print(f'opt_level={result["opt_level"]}:\nBaseline not trained: {result["not_trained"]}\nBaseline trained: {result["trained"]}\n')
-    # accelerator_not_trained, accelerator_trained, accelerator_outputs, accelerator_data = train_integration(zero_stage)
-    # assert (
-    #     baseline_not_trained["accuracy"] == accelerator_not_trained["accuracy"]
-    # ), f'ZERO stage {zero_stage}: Accuracy should be the same for the baseline and accelerator: {baseline_not_trained["accuracy"]} == {accelerator_not_trained["accuracy"]}'
-    # assert (
-    #     baseline_not_trained["f1"] == accelerator_not_trained["f1"]
-    # ), f'ZERO stage {zero_stage}: F1 score should be the same for the baseline and accelerator: {baseline_not_trained["f1"]} == {accelerator_not_trained["f1"]}'
-    # assert (
-    #     baseline_trained["accuracy"] == accelerator_trained["accuracy"]
-    # ), f'ZERO stage {zero_stage}: Accuracy should be the same for the baseline and accelerator: {baseline_trained["accuracy"]} == {accelerator_trained["accuracy"]}'
-    # assert (
-    #     baseline_trained["f1"] == accelerator_trained["f1"]
-    # ), f'ZERO stage {zero_stage}: F1 score should be the same for the baseline and accelerator: {baseline_trained["f1"]} == {accelerator_trained["f1"]}'
+    for zero_stage in [1, 2]:
+        for opt_level in ["O1", "O2", "O3"]:
+            baseline_not_trained, baseline_trained = train_baseline(zero_stage, opt_level)
+            accelerator_not_trained, accelerator_trained = train_integration(zero_stage, opt_level)
+            assert (
+                baseline_not_trained["accuracy"] == accelerator_not_trained["accuracy"]
+            ), f'ZERO stage {zero_stage}, opt_level={opt_level}:\nAccuracy should be the same for the baseline and accelerator: {baseline_not_trained["accuracy"]} == {accelerator_not_trained["accuracy"]}'
+            assert (
+                baseline_not_trained["f1"] == accelerator_not_trained["f1"]
+            ), f'ZERO stage {zero_stage}, opt_level={opt_level}:\nF1 score should be the same for the baseline and accelerator: {baseline_not_trained["f1"]} == {accelerator_not_trained["f1"]}'
+            assert (
+                baseline_trained["accuracy"] == accelerator_trained["accuracy"]
+            ), f'ZERO stage {zero_stage}, opt_level={opt_level}:\nAccuracy should be the same for the baseline and accelerator: {baseline_trained["accuracy"]} == {accelerator_trained["accuracy"]}'
+            assert (
+                baseline_trained["f1"] == accelerator_trained["f1"]
+            ), f'ZERO stage {zero_stage}, opt_level={opt_level}:\nF1 score should be the same for the baseline and accelerator: {baseline_trained["f1"]} == {accelerator_trained["f1"]}'
 
     torch.distributed.destroy_process_group()

benchmarks/fp8/ms_amp/fsdp.py

@@ -17,21 +17,21 @@
 
 This particular script verifies this for FSDP training.
 """
+from functools import partial
+
 import evaluate
 import msamp
 import torch
-from accelerate import FullyShardedDataParallelPlugin as FSDPPlugin
+from fp8_utils import evaluate_model, get_training_utilities
+from msamp.common.dtype import Dtypes
 from msamp.fsdp import FP8FullyShardedDataParallel
 from msamp.optim import FSDPAdamW
-from msamp.common.dtype import Dtypes
-from fp8_utils import evaluate_model, get_training_utilities
+from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
+from transformers.models.bert import BertLayer
 
 from accelerate import Accelerator
-from accelerate.state import AcceleratorState
+from accelerate import FullyShardedDataParallelPlugin as FSDPPlugin
 from accelerate.utils import FP8RecipeKwargs, set_seed
-from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
-from transformers.models.bert import BertLayer
-from functools import partial
 
 
 MODEL_NAME = "bert-base-cased"
@@ -45,10 +45,7 @@ def train_baseline(opt_level="O2"):
     accelerator = Accelerator()
     device = accelerator.device
     model, optimizer = msamp.initialize(
-        model, optimizer,
-        opt_level=opt_level,
-        weight_qtype=Dtypes.kfloat8_e4m3,
-        use_fsdp=True
+        model, optimizer, opt_level=opt_level, weight_qtype=Dtypes.kfloat8_e4m3, use_fsdp=True
     )
 
     model = FP8FullyShardedDataParallel(
@@ -60,7 +57,7 @@ def train_baseline(opt_level="O2"):
         backward_prefetch=None,
         forward_prefetch=False,
         limit_all_gathers=True,
-        device_id=device
+        device_id=device,
     )
     optimizer = FSDPAdamW(optimizer)
 
@@ -78,8 +75,8 @@ def train_baseline(opt_level="O2"):
 
     trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
 
-    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = (
-        accelerator.free_memory(model, optimizer, train_dataloader, eval_dataloader, lr_scheduler)
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.free_memory(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
     )
     assert (
         trained_model_results["accuracy"] > base_model_results["accuracy"]
@@ -118,8 +115,8 @@ def train_integration(opt_level="O2"):
 
     trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
 
-    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = (
-        accelerator.free_memory(model, optimizer, train_dataloader, eval_dataloader, lr_scheduler)
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.free_memory(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
     )
     assert (
         trained_model_results["accuracy"] > base_model_results["accuracy"]

src/accelerate/accelerator.py

@@ -104,7 +104,7 @@
     save_fsdp_optimizer,
     wait_for_everyone,
 )
-from .utils.constants import FSDP_PYTORCH_VERSION, PROFILE_PATTERN_NAME
+from .utils.constants import PROFILE_PATTERN_NAME
 from .utils.modeling import get_state_dict_offloaded_model
 from .utils.other import is_compiled_module
 
@@ -310,8 +310,8 @@ def __init__(
         # if os.environ.get("ACCELERATE_USE_FSDP", "false") == "true" or isinstance(
         #     fsdp_plugin, FullyShardedDataParallelPlugin
         # ):
-            # if is_torch_version("<", FSDP_PYTORCH_VERSION):
-            #     raise ValueError(f"FSDP requires PyTorch >= {FSDP_PYTORCH_VERSION}")
+        # if is_torch_version("<", FSDP_PYTORCH_VERSION):
+        #     raise ValueError(f"FSDP requires PyTorch >= {FSDP_PYTORCH_VERSION}")
 
         if fsdp_plugin is None:  # init from env variables
             fsdp_plugin = (
@@ -507,8 +507,11 @@ def __init__(
         elif self.state.mixed_precision == "fp8":
             # We always enable `native_amp` for FP8
             self.native_amp = True
-            # MS-AMP requires grad scaler however
-            if self.fp8_backend == "MSAMP" and self.distributed_type not in (DistributedType.FSDP, DistributedType.DEEPSPEED):
+            # MS-AMP requires `GradScaler` even with bf16 autocast w/ single GPU or DDP:
+            if self.fp8_backend == "MSAMP" and self.distributed_type not in (
+                DistributedType.FSDP,
+                DistributedType.DEEPSPEED,
+            ):
                 self.scaler = torch.cuda.amp.GradScaler()
 
         # Start of internal step tracking
@@ -1336,6 +1339,7 @@ def prepare(self, *args, device_placement=None):
             # We need to convert the underlying optimizer to FSDPAdamW *after* FSDP wrapping
             result = list(result)
             from msamp.optim import FSDPAdamW
+
             for i, obj in enumerate(result):
                 if isinstance(obj, AcceleratedOptimizer):
                     result[i].optimizer = FSDPAdamW(optimizer=obj.optimizer)
@@ -1636,6 +1640,13 @@ def _prepare_te(self, *args, device=None):
     def _prepare_deepspeed(self, *args):
         import deepspeed
 
+        ds_initialize = deepspeed.initialize
+        if self.fp8_backend == "MSAMP":
+            # MS-AMP requires DeepSpeed patches
+            from msamp import deepspeed as msamp_deepspeed
+
+            ds_initialize = msamp_deepspeed.initialize
+
         deepspeed_plugin = self.state.deepspeed_plugin
 
         is_dataloader_present = any(isinstance(obj, torch.utils.data.DataLoader) for obj in args)
@@ -1824,7 +1835,7 @@ def _prepare_deepspeed(self, *args):
                         if type(scheduler).__name__ in deepspeed.runtime.lr_schedules.VALID_LR_SCHEDULES:
                             kwargs["lr_scheduler"] = scheduler
 
-            engine, optimizer, _, lr_scheduler = deepspeed.initialize(**kwargs)
+            engine, optimizer, _, lr_scheduler = ds_initialize(**kwargs)
             if optimizer is not None:
                 optimizer = DeepSpeedOptimizerWrapper(optimizer)
             if scheduler is not None:
@@ -2037,8 +2048,10 @@ def _prepare_msamp(self, *args, device_placement):
                 # NOTE: MS-AMP fsdp relies on it's own MP policy, we must drop the users
                 self.state.fsdp_plugin.mixed_precision_policy = None
             from msamp.common.dtype import Dtypes
+
             model, optimizer = msamp.initialize(
-                model, optimizer,
+                model,
+                optimizer,
                 opt_level=self.fp8_recipe_handler.opt_level,
                 use_fsdp=self.distributed_type == DistributedType.FSDP,
                 weight_qtype=Dtypes.kfloat8_e4m3,
@@ -3595,4 +3608,6 @@ def fp8_backend(self):
         "Returns the configured backend for training in FP8"
         if self.mixed_precision == "fp8" and self.fp8_recipe_handler is not None:
             return self.fp8_recipe_handler.backend
+        elif self.state.deepspeed_plugin is not None and self.state.deepspeed_plugin.enable_msamp:
+            return "MSAMP"
         return None

src/accelerate/utils/dataclasses.py

@@ -972,6 +972,16 @@ class DeepSpeedPlugin:
             " `MixtralSparseMoeBlock`, `Qwen2MoeSparseMoeBlock`, `JetMoEAttention,JetMoEBlock` ..."
         },
     )
+    enable_msamp: bool = field(
+        default=None,
+        metadata={"help": "Flag to indicate whether to enable MS-AMP backend for FP8 training."},
+    )
+    msamp_opt_level: str = field(
+        default=None,
+        metadata={
+            "help": "Optimization level for MS-AMP. Only applicable if `enable_msamp` is True. Should be one of ['O1', 'O2', 'O3']."
+        },
+    )
 
     def __post_init__(self):
         from .deepspeed import HfDeepSpeedConfig
@@ -1006,6 +1016,12 @@ def __post_init__(self):
                 os.environ.get("ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL", "false") == "true"
             )
 
+        if self.enable_msamp is None:
+            self.enable_msamp = os.environ.get("ACCELERATE_FP8_BACKEND", None) == "MSAMP"
+
+        if self.msamp_opt_level is None:
+            self.msamp_opt_level = os.environ.get("ACCELERATE_FP8_OPT_LEVEL", "O1")
+
         if self.hf_ds_config is None:
             self.hf_ds_config = os.environ.get("ACCELERATE_DEEPSPEED_CONFIG_FILE", "none")
         if (
@@ -1075,6 +1091,14 @@ def __post_init__(self):
         if self.zero3_init_flag and not self.hf_ds_config.is_zero3():
             warnings.warn("DeepSpeed Zero3 Init flag is only applicable for ZeRO Stage 3. Setting it to False.")
             self.zero3_init_flag = False
+        if self.enable_msamp:
+            if self.zero_stage == 3:
+                raise NotImplementedError(
+                    "MS-AMP is not supported for ZeRO Stage 3. Please use ZeRO Stage 0, 1, or 2 instead."
+                )
+            if self.msamp_opt_level not in ["O1", "O2", "O3"]:
+                raise ValueError("Invalid optimization level for MS-AMP. Please use one of ['O1', 'O2', 'O3'].")
+            self.deepspeed_config["msamp"] = {"enabled": True, "opt_level": self.msamp_opt_level}
 
     def fill_match(self, ds_key_long, mismatches=None, must_match=True, **kwargs):
         mismatches = [] if mismatches is None else mismatches
@@ -1144,6 +1168,10 @@ def set_mixed_precision(self, mixed_precision):
             if "bf16" not in ds_config:
                 ds_config["bf16"] = {"enabled": True}
 
+        if mixed_precision == "fp8" and self.enable_msamp:
+            if "msamp" not in ds_config:
+                ds_config["msamp"] = {"enabled": True, "opt_level": self.msamp_opt_level}
+
         if mixed_precision != "no":
             diff_dtype = "bf16" if mixed_precision == "fp16" else "fp16"
             if str(ds_config.get(diff_dtype, {}).get("enabled", "False")).lower() == "true":