Upstream Accelerate

.github/workflows/test_accelerate_gaudi2.yml

@@ -0,0 +1,37 @@
+name: Accelerate integration tests
+
+on:
+  workflow_dispatch:
+  pull_request:
+    branches: [main]
+  push:
+    branches: [main]
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+jobs:
+  accelerate:
+    name: Test Accelerate integration
+    runs-on: [self-hosted, linux, x64, gaudi2, fast]
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+      - name: Pull image
+        run: |
+          docker pull vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest
+      - name: Run tests
+        run: |
+          docker run \
+          --rm \
+          --ipc=host \
+          --runtime=habana \
+          --cap-add=sys_nice \
+          -v $PWD:/root/workspace \
+          --workdir=/root/workspace \
+          -e OMPI_MCA_btl_vader_single_copy_mechanism=none \
+          -e HABANA_VISIBLE_DEVICES=$DOCKER_HABANA_VISIBLE_DEVICES \
+          vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest \
+          /bin/bash tests/ci/accelerate.sh

examples/stable-diffusion/training/textual_inversion.py

@@ -32,6 +32,7 @@
 import torch.nn.functional as F
 import torch.utils.checkpoint
 import transformers
+from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import ProjectConfiguration
 from diffusers import (
@@ -51,7 +52,6 @@
 from transformers import CLIPTextModel, CLIPTokenizer
 
 from optimum.habana import GaudiConfig
-from optimum.habana.accelerate import GaudiAccelerator
 from optimum.habana.diffusers import GaudiDDIMScheduler, GaudiStableDiffusionPipeline
 from optimum.habana.utils import set_seed
 
@@ -588,12 +588,11 @@ def main():
 
     gaudi_config = GaudiConfig.from_pretrained(args.gaudi_config_name)
 
-    accelerator = GaudiAccelerator(
+    accelerator = Accelerator(
         gradient_accumulation_steps=args.gradient_accumulation_steps,
         mixed_precision="bf16" if gaudi_config.use_torch_autocast or args.bf16 else "no",
         log_with=args.report_to,
         project_config=accelerator_project_config,
-        force_autocast=gaudi_config.use_torch_autocast or args.bf16,
     )
 
     if args.report_to == "wandb":

examples/stable-diffusion/training/textual_inversion_sdxl.py

@@ -31,6 +31,7 @@
 import torch.nn.functional as F
 import torch.utils.checkpoint
 import transformers
+from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import ProjectConfiguration
 from diffusers import (
@@ -48,7 +49,6 @@
 from tqdm.auto import tqdm
 
 from optimum.habana import GaudiConfig
-from optimum.habana.accelerate import GaudiAccelerator
 from optimum.habana.diffusers import (
     GaudiStableDiffusionXLPipeline,
 )
@@ -576,12 +576,11 @@ def main():
 
     gaudi_config = GaudiConfig.from_pretrained(args.gaudi_config_name)
 
-    accelerator = GaudiAccelerator(
+    accelerator = Accelerator(
         gradient_accumulation_steps=args.gradient_accumulation_steps,
         mixed_precision="bf16" if gaudi_config.use_torch_autocast or args.bf16 else "no",
         log_with=args.report_to,
         project_config=accelerator_project_config,
-        force_autocast=gaudi_config.use_torch_autocast or args.bf16,
     )
 
     if args.report_to == "wandb":

examples/stable-diffusion/training/train_controlnet.py

@@ -35,6 +35,7 @@
 import torch.nn.functional as F
 import torch.utils.checkpoint
 import transformers
+from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import ProjectConfiguration
 from datasets import load_dataset
@@ -54,7 +55,6 @@
 from transformers import AutoTokenizer, PretrainedConfig
 
 from optimum.habana import GaudiConfig
-from optimum.habana.accelerate import GaudiAccelerator
 from optimum.habana.diffusers import GaudiDDIMScheduler, GaudiStableDiffusionControlNetPipeline
 from optimum.habana.utils import set_seed
 
@@ -765,12 +765,11 @@ def main(args):
     # Set autocast to True for --bf16
     if args.bf16:
         gaudi_config.use_torch_autocast = True
-    accelerator = GaudiAccelerator(
+    accelerator = Accelerator(
         gradient_accumulation_steps=args.gradient_accumulation_steps,
         mixed_precision="bf16" if gaudi_config.use_torch_autocast else "no",
         log_with=args.report_to,
         project_config=accelerator_project_config,
-        force_autocast=gaudi_config.use_torch_autocast,
     )
 
     # Make one log on every process with the configuration for debugging.

examples/stable-diffusion/training/train_dreambooth.py

@@ -39,8 +39,10 @@
 import torch.nn.functional as F
 import torch.utils.checkpoint
 import transformers
+from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import DistributedDataParallelKwargs
+from accelerate.utils.dataclasses import DistributedType
 from diffusers import (
     AutoencoderKL,
     DDPMScheduler,
@@ -60,8 +62,6 @@
 from transformers import AutoTokenizer, PretrainedConfig
 
 from optimum.habana import GaudiConfig
-from optimum.habana.accelerate import GaudiAccelerator
-from optimum.habana.accelerate.utils.dataclasses import GaudiDistributedType
 from optimum.habana.diffusers import GaudiStableDiffusionPipeline
 from optimum.habana.transformers.trainer import _is_peft_model
 from optimum.habana.utils import set_seed
@@ -834,12 +834,11 @@ def main(args):
 
     gaudi_config = GaudiConfig.from_pretrained(args.gaudi_config_name)
     gaudi_config.use_torch_autocast = gaudi_config.use_torch_autocast or args.mixed_precision == "bf16"
-    accelerator = GaudiAccelerator(
+    accelerator = Accelerator(
         gradient_accumulation_steps=args.gradient_accumulation_steps,
         mixed_precision=args.mixed_precision,
         log_with=args.report_to,
         project_dir=logging_dir,
-        force_autocast=gaudi_config.use_torch_autocast,
     )
     if args.report_to == "wandb":
         import wandb
@@ -1088,7 +1087,7 @@ def unwrap_model(model, training=False):
         if not training:
             return model
         else:
-            if accelerator.distributed_type == GaudiDistributedType.MULTI_HPU:
+            if accelerator.distributed_type == DistributedType.MULTI_HPU:
                 kwargs = {}
                 kwargs["gradient_as_bucket_view"] = True
                 accelerator.ddp_handler = DistributedDataParallelKwargs(**kwargs)

examples/stable-diffusion/training/train_dreambooth_lora_flux.py

@@ -34,8 +34,10 @@
 import torch
 import torch.utils.checkpoint
 import transformers
+from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import DistributedDataParallelKwargs, ProjectConfiguration
+from accelerate.utils.dataclasses import DistributedType
 from datasets import load_dataset
 from diffusers import (
     AutoencoderKL,
@@ -68,8 +70,6 @@
 from transformers import T5EncoderModel
 
 from optimum.habana import GaudiConfig
-from optimum.habana.accelerate import GaudiAccelerator
-from optimum.habana.accelerate.utils.dataclasses import GaudiDistributedType
 from optimum.habana.utils import set_seed
 
 
@@ -643,12 +643,11 @@ def main(args):
 
     accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
     kwargs = DistributedDataParallelKwargs(find_unused_parameters=False)
-    accelerator = GaudiAccelerator(
+    accelerator = Accelerator(
         gradient_accumulation_steps=args.gradient_accumulation_steps,
         mixed_precision=args.mixed_precision,
         log_with=args.report_to,
         project_config=accelerator_project_config,
-        force_autocast=gaudi_config.use_torch_autocast,
         kwargs_handlers=[kwargs],
     )
 
@@ -762,7 +761,7 @@ def save_model_hook(models, weights, output_dir):
     def load_model_hook(models, input_dir):
         transformer_ = None
 
-        if not accelerator.distributed_type == GaudiDistributedType.DEEPSPEED:
+        if not accelerator.distributed_type == DistributedType.DEEPSPEED:
             while len(models) > 0:
                 model = models.pop()
 
@@ -1075,7 +1074,7 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
                 progress_bar.update(1)
                 global_step += 1
 
-                if accelerator.is_main_process or accelerator.distributed_type == GaudiDistributedType.DEEPSPEED:
+                if accelerator.is_main_process or accelerator.distributed_type == DistributedType.DEEPSPEED:
                     if global_step % args.checkpointing_steps == 0:
                         # _before_ saving state, check if this save would set us over the `checkpoints_total_limit`
                         if args.checkpoints_total_limit is not None:

examples/stable-diffusion/training/train_dreambooth_lora_sdxl.py

@@ -35,8 +35,10 @@
 import torch.nn.functional as F
 import torch.utils.checkpoint
 import transformers
+from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import DistributedDataParallelKwargs
+from accelerate.utils.dataclasses import DistributedType
 from diffusers import (
     AutoencoderKL,
     DDPMScheduler,
@@ -67,8 +69,6 @@
 from transformers import AutoTokenizer, PretrainedConfig
 
 from optimum.habana import GaudiConfig
-from optimum.habana.accelerate import GaudiAccelerator
-from optimum.habana.accelerate.utils.dataclasses import GaudiDistributedType
 from optimum.habana.diffusers import GaudiStableDiffusionXLPipeline
 from optimum.habana.transformers.trainer import _is_peft_model
 from optimum.habana.utils import set_seed
@@ -821,12 +821,11 @@ def main(args):
     logging_dir = Path(args.output_dir, args.logging_dir)
     gaudi_config = GaudiConfig.from_pretrained(args.gaudi_config_name)
     gaudi_config.use_torch_autocast = gaudi_config.use_torch_autocast or args.mixed_precision == "bf16"
-    accelerator = GaudiAccelerator(
+    accelerator = Accelerator(
         gradient_accumulation_steps=args.gradient_accumulation_steps,
         mixed_precision=args.mixed_precision,
         log_with=args.report_to,
         project_dir=logging_dir,
-        force_autocast=gaudi_config.use_torch_autocast,
     )
     if args.report_to == "wandb":
         if not is_wandb_available():
@@ -1019,7 +1018,7 @@ def unwrap_model(model, training=False):
         if not training:
             return model
         else:
-            if accelerator.distributed_type == GaudiDistributedType.MULTI_HPU:
+            if accelerator.distributed_type == DistributedType.MULTI_HPU:
                 kwargs = {}
                 kwargs["gradient_as_bucket_view"] = True
                 accelerator.ddp_handler = DistributedDataParallelKwargs(**kwargs)

examples/stable-diffusion/training/train_text_to_image_sdxl.py

@@ -41,8 +41,10 @@
 import torch.nn.functional as F
 import torch.utils.checkpoint
 import transformers
+from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import DistributedDataParallelKwargs, ProjectConfiguration
+from accelerate.utils.dataclasses import DistributedType
 from datasets import load_dataset
 from diffusers import (
     AutoencoderKL,
@@ -61,8 +63,6 @@
 from transformers import AutoTokenizer, PretrainedConfig
 
 from optimum.habana import GaudiConfig
-from optimum.habana.accelerate import GaudiAccelerator
-from optimum.habana.accelerate.utils.dataclasses import GaudiDistributedType
 from optimum.habana.diffusers import (
     GaudiDDIMScheduler,
     GaudiEulerAncestralDiscreteScheduler,
@@ -714,12 +714,11 @@ def main(args):
 
     gaudi_config = GaudiConfig.from_pretrained(args.gaudi_config_name)
     gaudi_config.use_torch_autocast = gaudi_config.use_torch_autocast or args.bf16
-    accelerator = GaudiAccelerator(
+    accelerator = Accelerator(
         gradient_accumulation_steps=args.gradient_accumulation_steps,
         mixed_precision="bf16" if gaudi_config.use_torch_autocast else "no",
         log_with=args.report_to,
         project_config=accelerator_project_config,
-        force_autocast=gaudi_config.use_torch_autocast,
     )
 
     if args.report_to == "wandb":
@@ -896,7 +895,7 @@ def main(args):
                         for idx, dt in enumerate(dataset["train"]):
                             dt["image"].save(f"{args.mediapipe}/{idx}.jpg")
                             f.write(dt["text"] + "\n")
-            if accelerator.distributed_type != GaudiDistributedType.NO:
+            if accelerator.distributed_type != DistributedType.NO:
                 torch.distributed.barrier()
 
             from media_pipe_imgdir import get_dataset_for_pipeline
@@ -1145,7 +1144,7 @@ def unwrap_model(model, training=False):
         if not training:
             return model
         else:
-            if accelerator.distributed_type == GaudiDistributedType.MULTI_HPU:
+            if accelerator.distributed_type == DistributedType.MULTI_HPU:
                 kwargs = {}
                 kwargs["gradient_as_bucket_view"] = True
                 accelerator.ddp_handler = DistributedDataParallelKwargs(**kwargs)

examples/trl/ppo.py

@@ -5,14 +5,14 @@
 from typing import List, Optional
 
 import torch
+from accelerate import Accelerator
 from datasets import load_dataset
 from peft import LoraConfig
 from tqdm import tqdm
 from transformers import Adafactor, AutoModelForSequenceClassification, AutoTokenizer, HfArgumentParser, pipeline
 from trl import AutoModelForCausalLMWithValueHead
 from trl.core import LengthSampler
 
-from optimum.habana.accelerate import GaudiAccelerator
 from optimum.habana.trl import GaudiPPOConfig, GaudiPPOTrainer, adapt_PreTrainedModelWrapper_to_gaudi
 from optimum.habana.utils import set_seed
 
@@ -191,7 +191,7 @@ def collator(data):
 set_seed(config.seed)
 
 # Now let's build the model, the reference model, and the tokenizer.
-current_device = GaudiAccelerator().local_process_index
+current_device = Accelerator().local_process_index
 lora_config = LoraConfig(
     r=script_args.lora_r,
     lora_alpha=script_args.lora_alpha,

optimum/habana/transformers/modeling_utils.py

@@ -16,8 +16,8 @@
 import transformers
 import transformers.utils.fx
 
-from ..accelerate.utils import extract_model_from_parallel
-from ..accelerate.utils.modeling import gaudi_check_device_same
+from ..local_accelerate.utils import extract_model_from_parallel
+from ..local_accelerate.utils.modeling import gaudi_check_device_same
 from ..quantizers.bitsandbytes import (
     gaudi_bitsandbytesconfig_post_init,
     gaudi_create_quantized_param,