use FusedAdam from TE

Signed-off-by: Peter St. John <[email protected]>

Author: pstjohn

bionemo-recipes/recipes/esm2_native_te/hydra_config/defaults.yaml

@@ -53,7 +53,7 @@ fp8_config:
 # Optimizer config
 adamw_kwargs:
   lr: 4e-4
-  fused: true
+  adam_w_mode: true
   betas: [0.9, 0.98]
   eps: 1e-8
   weight_decay: 0.01

bionemo-recipes/recipes/esm2_native_te/train_ddp.py

@@ -21,8 +21,8 @@
 import transformer_engine.pytorch
 from omegaconf import DictConfig
 from torch.distributed.device_mesh import init_device_mesh
-from torch.optim import AdamW
 from transformer_engine.common.recipe import Format
+from transformer_engine.pytorch.optimizers import FusedAdam
 from transformers import AutoConfig, AutoModelForMaskedLM
 
 from checkpoint import load_checkpoint_ddp, save_checkpoint_ddp, save_final_model_ddp, should_save_checkpoint
@@ -81,7 +81,7 @@ def main(args: DictConfig) -> float | None:
         pass
 
     # Create optimizer.
-    optimizer = AdamW(model.parameters(), **args.adamw_kwargs)
+    optimizer = FusedAdam(model.parameters(), **args.adamw_kwargs)
     scheduler = get_linear_schedule_with_warmup(optimizer, **args.lr_scheduler_kwargs)
 
     model = model.to(device=device)

bionemo-recipes/recipes/esm2_native_te/train_fsdp2.py

@@ -22,8 +22,8 @@
 from omegaconf import DictConfig, OmegaConf
 from torch.distributed.device_mesh import init_device_mesh
 from torch.distributed.fsdp import fully_shard
-from torch.optim import AdamW
 from transformer_engine.common.recipe import Format
+from transformer_engine.pytorch.optimizers import FusedAdam
 from transformers import AutoConfig, AutoModelForMaskedLM
 
 # This import seems to be needed with meta device init and AutoModel.from_config
@@ -87,7 +87,7 @@ def main(args: DictConfig) -> float | None:  # noqa: C901
     fully_shard(model, mesh=device_mesh["dp"])
 
     # Create optimizer. Convert OmegaConf to regular dict to avoid serialization issues (BIONEMO-2873).
-    optimizer = AdamW(model.parameters(), **OmegaConf.to_container(args.adamw_kwargs, resolve=True))  # type: ignore
+    optimizer = FusedAdam(model.parameters(), **OmegaConf.to_container(args.adamw_kwargs, resolve=True))  # type: ignore
     scheduler = get_linear_schedule_with_warmup(optimizer, **args.lr_scheduler_kwargs)
 
     if args.use_meta_device:

bionemo-recipes/recipes/esm2_native_te/train_mfsdp.py

@@ -23,8 +23,8 @@
 from megatron_fsdp.fully_shard import fully_shard
 from omegaconf import DictConfig, OmegaConf
 from torch.distributed.device_mesh import init_device_mesh
-from torch.optim import AdamW
 from transformer_engine.common.recipe import Format
+from transformer_engine.pytorch.optimizers import FusedAdam
 from transformers import AutoConfig, AutoModelForMaskedLM
 
 from checkpoint import load_checkpoint_mfsdp, save_checkpoint_mfsdp, save_final_model_mfsdp, should_save_checkpoint
@@ -85,7 +85,7 @@ def main(args: DictConfig) -> float | None:
     logger.info("Initialized Model:\n%s", model)
 
     # Create optimizer. Convert OmegaConf to regular dict to avoid serialization issues (BIONEMO-2873).
-    optimizer = AdamW(model.parameters(), **OmegaConf.to_container(args.adamw_kwargs, resolve=True))  # type: ignore
+    optimizer = FusedAdam(model.parameters(), **OmegaConf.to_container(args.adamw_kwargs, resolve=True))  # type: ignore
 
     # Wrap model in megatron-fsdp
     model, optimizer = fully_shard(