NVIDIA-NeMo / NVIDIA-NeMo/RL

Dtensor policy v1 and v2 future plan

Open
#1,566 0 comments 0 reactions 0 assignees View on GitHub
enhancement t-pytdensor
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

There are several issues about dtensor policy v1 and v2 now.
1. There are problems for v2 (NeMo Automodel to support some diffusion models.)
2. Researchers enjoy v1 merits: fully HF native and transparency.

But there are other issues from SW perspective:
1. NeMo Automodel will be much more powerful and feature rich built on top of FSDP.
2. There are many duplicated codes in v1 and v2 which will make maintenance more difficult.

With many discussions between researching we had some conclusions:
1. Keep v1 till they feel comfortable to deprecate.
2. Increase Automodel support for gapped features.
3. Automodel will have HF native fallback mode and this will be supported with in v2 dtensor policy worker.
4. Our refactor works on v2 policy worker will also apply to v1 to make code cleaner on both side.
5. There will be divergent features between v1 and v2 in future releases including EP, PP, HSDP, CP + Seq pack via TE.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.