updating esm2 native recipe (#1078)

* Adds separate `train_ddp.py`, `train_fsdp2.py`, and `train_nvfsdp.py`
entrypoints
* Adds comparison against FA-2 based HF transformers model

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

* **New Features**
* Distributed ESM‑2 training entrypoints (DDP, MFSDP/FSDP2) and a shared
linear warmup/decay LR scheduler.

* **Configuration**
* Switched defaults to nvFSDP-style sharding; updated model identifiers
and training hyperparameters (train steps, warmup, optimizer LR).

* **Documentation**
  * Added new ESM‑2 training README; removed an outdated README.

* **Build/Chores**
* Install Transformers from the Git repo; Docker builds can use netrc
credentials for installs.

* **Tests**
* Added extensive single‑ and multi‑GPU training tests; removed obsolete
tests.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

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

Author: pstjohn

.devcontainer/recipes/requirements.txt

@@ -7,6 +7,6 @@ megatron-fsdp==0.1.0rc0
 torchmetrics
 tqdm
 transformer_engine
-transformers
+transformers @ git+https://github.com/huggingface/transformers.git
 typer
 wandb

recipes/README.md

@@ -135,9 +135,7 @@ import torch
 from torch.distributed import init_process_group, destroy_process_group
 
 
-@hydra.main(
-    config_path="hydra_config", config_name="L0_sanity.yaml", version_base="1.2"
-)
+@hydra.main(config_path="hydra_config", config_name="L0_sanity", version_base="1.2")
 def main(args: DictConfig):
     """Main training entrypoint."""
 
@@ -306,7 +304,7 @@ def test_accelerate_launch(accelerate_config, tmp_path):
             str(accelerate_config_path),
             "train.py",
             "--config-name",
-            "L0_sanity.yaml",
+            "L0_sanity",
             f"trainer.output_dir={tmp_path}",
         ],
         cwd=recipe_dir,

recipes/amplify_accelerate_te_fp8/test_train.py

@@ -187,7 +187,7 @@ def test_accelerate_launch(accelerate_config, tmp_path):
             str(accelerate_config_path),
             str(train_py),
             "--config-name",
-            "L0_sanity.yaml",
+            "L0_sanity",
             f"trainer.output_dir={tmp_path}",
         ],
         cwd=recipe_dir,

recipes/amplify_accelerate_te_fp8/train.py

@@ -32,7 +32,7 @@
 logger = logging.getLogger(__name__)
 
 
-@hydra.main(config_path="hydra_config", config_name="L0_sanity.yaml", version_base="1.2")
+@hydra.main(config_path="hydra_config", config_name="L0_sanity", version_base="1.2")
 def main(args: DictConfig):
     """Entrypoint."""
     config = AutoConfig.from_pretrained(args.model_tag, trust_remote_code=True)

…ipes/esm2_native_te_nvfsdp/.dockerignore …cipes/esm2_native_te_mfsdp/.dockerignorerecipes/esm2_native_te_nvfsdp/.dockerignore renamed to recipes/esm2_native_te_mfsdp/.dockerignore recipes/esm2_native_te_nvfsdp/.dockerignore renamed to recipes/esm2_native_te_mfsdp/.dockerignore

@@ -1,7 +1,8 @@
 # syntax=docker/dockerfile:1.4
 FROM nvcr.io/nvidia/pytorch:25.06-py3
 
-RUN --mount=type=cache,target=/root/.cache/pip \
+RUN --mount=type=secret,id=netrc,target=/root/.netrc \
+    --mount=type=cache,target=/root/.cache/pip \
     --mount=type=bind,source=requirements.txt,target=/requirements.txt \
     PIP_CONSTRAINT= pip install -r /requirements.txt
 

recipes/esm2_native_te_nvfsdp/.ruff.toml recipes/esm2_native_te_mfsdp/.ruff.tomlrecipes/esm2_native_te_nvfsdp/.ruff.toml renamed to recipes/esm2_native_te_mfsdp/.ruff.toml

@@ -0,0 +1,15 @@
+# ESM-2 training with megatron-fsdp and custom pytorch training loop
+
+Build the docker image with the following command:
+
+```bash
+docker build -t my_image .
+```
+
+## Running training
+
+Run training with
+
+```bash
+docker run --rm -it --gpus all my_image torchrun train_mfsdp.py --config-name L0_sanity
+```

recipes/esm2_native_te_nvfsdp/Dockerfile recipes/esm2_native_te_mfsdp/Dockerfilerecipes/esm2_native_te_nvfsdp/Dockerfile renamed to recipes/esm2_native_te_mfsdp/Dockerfile

@@ -2,9 +2,9 @@ defaults:
   - defaults
 
 # Training config
-model_name: esm2_t6_8M_UR50D
+model_name: nvidia/esm2_t6_8M_UR50D
 micro_batch_size: 2
-num_train_steps: 5
+num_train_steps: 250
 
 # WandB config
 wandb_init_args:
@@ -13,5 +13,7 @@ wandb_init_args:
 
 # Learning rate scheduler config
 lr_scheduler_kwargs:
-  num_warmup_steps: 2
-  num_training_steps: 8
+  num_warmup_steps: 0
+
+adamw_kwargs:
+  lr: 1e-2