Refactor accelerate test suite, add amplify tests

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

Author: pstjohn

recipes/esm2_accelerate/accelerate_config/dynamo.yaml

@@ -0,0 +1,18 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: MULTI_GPU
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 1
+rdzv_backend: c10d
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false
+dynamo_config:
+  dynamo_backend: INDUCTOR

recipes/esm2_accelerate/accelerate_config/fp8.yaml

@@ -0,0 +1,28 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: MULTI_GPU
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+machine_rank: 0
+main_training_function: main
+mixed_precision: fp8
+fp8_config:
+  amax_compute_algorithm: max
+  amax_history_length: 1024
+  backend: TE
+  fp8_format: HYBRID
+  interval: 1
+  margin: 0
+  override_linear_precision:
+  - false
+  - false
+  - false
+  use_autocast_during_eval: false
+num_machines: 1
+num_processes: 1
+rdzv_backend: c10d
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

recipes/esm2_accelerate/hydra_config/L0_sanity.yaml

@@ -13,5 +13,5 @@ trainer:
   eval_steps: 1000
   logging_steps: 10
   report_to: "none"
-  dataloader_num_workers: 0
+  dataloader_num_workers: 4
   warmup_steps: 0

recipes/esm2_accelerate/hydra_config/L0_sanity_amplify.yaml

@@ -0,0 +1,16 @@
+defaults:
+  - defaults
+  - _self_
+
+model_tag: "nvidia/AMPLIFY_120M"
+stop_after_n_steps: 250
+
+trainer:
+  run_name: "amplify_120M_sanity"
+  per_device_train_batch_size: 2
+  per_device_eval_batch_size: 2
+  save_steps: 1000
+  eval_steps: 1000
+  logging_steps: 10
+  report_to: "none"
+  dataloader_num_workers: 4

recipes/esm2_accelerate/hydra_config/defaults_amplify.yaml

@@ -0,0 +1,16 @@
+defaults:
+  - defaults
+  - _self_
+
+model_tag: "nvidia/AMPLIFY_350M"
+stop_after_n_steps: 1_000_000
+max_seq_length: 512
+
+trainer:
+  learning_rate: 1e-3
+  adam_beta2: 0.95
+  lr_scheduler_type: "cosine_with_min_lr"
+  lr_scheduler_kwargs:
+    min_lr: 1e-4
+  warmup_steps: 1_000
+  max_steps: 1_000_000