make sure we initialize accelerator before model (#1132)

We need to initialize the `Accelerator` object before creating TE layers
or they all end up on a single device

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- New Features
- Added a ready-to-run performance test preset for the esm2 t48 15B
model with sensible defaults: model tag, step cap, batch sizes, learning
rate, weight decay, warmup steps, and Weights & Biases logging.

- Bug Fixes
- Improved multi-GPU initialization by starting distributed state
earlier, reducing setup issues and OOM risk without changing training
behavior.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

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 PartialState
 from omegaconf import DictConfig
 from transformers import AutoConfig, AutoModelForMaskedLM
 from transformers.trainer import Trainer
@@ -35,6 +36,15 @@
 @hydra.main(config_path="hydra_config", config_name="L0_sanity", version_base="1.2")
 def main(args: DictConfig):
     """Entrypoint."""
+    # Initialize Accelerate's distributed state early so torch device is set per process
+    state = PartialState()
+    logger.info(
+        "Accelerate initialized (local_process_index=%s, num_processes=%s, device=%s)",
+        state.local_process_index,
+        state.num_processes,
+        state.device,
+    )
+
     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 +67,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)