make sure we initialize accelerator before model

Signed-off-by: Peter St. John <pstjohn@nvidia.com>

Author: pstjohn

recipes/esm2_accelerate/hydra_config/L1_15B_perf_test.yaml

@@ -0,0 +1,14 @@
+defaults:
+  - defaults
+  - _self_
+
+model_tag: nvidia/esm2_t48_15B_UR50D
+stop_after_n_steps: 500
+trainer:
+  run_name: "esm2_t48_15B_UR50D_perf"
+  per_device_train_batch_size: 12
+  per_device_eval_batch_size: 12
+  report_to: "wandb"
+  learning_rate: 1.6e-4
+  weight_decay: 0.1
+  warmup_steps: 20_000

recipes/esm2_accelerate/train.py

@@ -19,6 +19,7 @@
 import hydra
 import torch
 import transformers
+from accelerate import Accelerator
 from omegaconf import DictConfig
 from transformers import AutoConfig, AutoModelForMaskedLM
 from transformers.trainer import Trainer
@@ -35,6 +36,10 @@
 @hydra.main(config_path="hydra_config", config_name="L0_sanity", version_base="1.2")
 def main(args: DictConfig):
     """Entrypoint."""
+    # We need to initialize the Accelerator manually prior to creating our model, otherwise we won't end up setting the
+    # current torch device and the model creation will all happen on a single GPU, typically leading to an OOM.
+    _ = Accelerator()
+
     config = AutoConfig.from_pretrained(args.model_tag, trust_remote_code=True)
     config.max_seq_length = args.max_seq_length
     config.micro_batch_size = args.trainer.per_device_train_batch_size
@@ -57,8 +62,6 @@ def main(args: DictConfig):
         callbacks=[StopAfterNStepsCallback(args.stop_after_n_steps)],
     )
 
-    logger.info("ACCELERATE STATE:\n%s\n", trainer.accelerator.state)
-
     if training_args.do_train:
         Path(training_args.output_dir).mkdir(parents=True, exist_ok=True)
         last_checkpoint = transformers.trainer_utils.get_last_checkpoint(training_args.output_dir)