Skip to content

Add Transformer-Engine Fused_Adam Optimizer Support #2293

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
vivekgoe wants to merge 2 commits into
base: main
Choose a base branch
from
+17 −0
Draft

Add Transformer-Engine Fused_Adam Optimizer Support #2293

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4d61ff7
Add TE_FusedAdamW optimizer support in build_optimizers function
vivekgoe Jan 28, 2026
fd87150
Fix typo in importing TE Fused_Adam
vivekgoe Jan 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions torchtitan/components/optimizer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -311,6 +311,23 @@ def build_optimizers(
"Adam": torch.optim.Adam,
"AdamW": torch.optim.AdamW,
}

if name == 'TE_FusedAdamW':
try:
from transformer_engine.pytorch.optimizers.fused_adam import FusedAdam
except (ImportError, ModuleNotFoundError) as e:
raise ImportError(
"FusedAdam optimizer could not be imported from transformer_engine. "
"Install transformer_engine with: "
"NVTE_FRAMEWORK=pytorch pip install --no-build-isolation "
"git+https://github.com/NVIDIA/TransformerEngine.git@stable "
"or use another optimizer"
) from e
optimizer_classes["TE_FusedAdamW"] = FusedAdam
del optimizer_kwargs['fused'], optimizer_kwargs['foreach']
optimizer_kwargs['exp_avg_dtype'] = torch.bfloat16
optimizer_kwargs['exp_avg_sq_dtype'] = torch.bfloat16
Comment on lines +328 to +329
Copy link
Contributor

@tianyu-l tianyu-l Feb 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @vivekgoe

My impression is that low-precision optimizer states and TE dependency are two separate topics, each worth its own discussions.

For low-precision optimizer states

  • @vivekgoe Could you show evidences that this is becoming the default for training?
  • @janeyx99 Does pytorch optimizer plan to support this feature?

In terms of dependency, I know that TE is an important and popular optimization library. However, torchtitan has been about pytorch native, prioritizing simplicity & maintainability, and platform neutral.

  • In the past, we have asked other hardware partners to create their own fork to host special optimizations. We try to organize the codebase to support extensibility (while not losing too much readbiility) so that modification on top is simple.
  • On the other hand, we recommend shipping critical features to pytorch before they are integrated in torchtitan.

Would love to hear your thoughts on this.

Copy link
Contributor Author

@vivekgoe vivekgoe Feb 4, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tianyu-l Thanks for detailed review comments. I agree It makes sense to separate out adopting low-precision optimizer states from creating TE dependency.

Regarding your question "Could you show evidences that this is becoming the default for training?", as far as I know DeepSeek pioneered this approach and I have not seen this used elsewhere, perhaps because support for this is not available in off the shelf optimizers.
I did a quick check on Qwen3-32B and LLama4-17Bx16e to confirm if models other than DeepSeek maintain quality (along with smaller memory footprint) with this feature, results look promising (see plots below). I can create a issue in pytorch repo to check if there is interest in adding this feature to torch.optim.adamw directly.
Code_Generated_Image (2)
Code_Generated_Image (3)

Regarding adding TE dependency, you make very valid arguments. As a short to medium term solution, will it be ok to move TE dependency to a new folder within experiment area so that users who wish to benefit from this feature can use it? Longer term we can add feature to native pytorch optimizer (assuming maintainers agree) and remove it from experiment area.

Copy link
Contributor

@tianyu-l tianyu-l Feb 7, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As a short to medium term solution, will it be ok to move TE dependency to a new folder within experiment area so that users who wish to benefit from this feature can use it? Longer term we can add feature to native pytorch optimizer (assuming maintainers agree) and remove it from experiment area.

The long-term plan sounds good to me.

For the short-term plan, the feedback we heard is that the current experiments folder is too messy, and we plan to clean things up (by deleting as much as possible). That doesn't mean we won't consider TE as an experiment, but we would like to see a clear vision and integration/maintenance plan, so that we can discuss and review together if it makes sense to the community. If it's just a feature like this fused adam optimizer, maybe showcasing it in a PR is good enough.

Hope it is not nonsense to you.


if name not in optimizer_classes:
raise NotImplementedError(f"Optimizer {name} not added.")
optimizer_cls = optimizer_classes[name]
Expand Down