Skip to content

Commit

Permalink
Romeyn/sampler (NVIDIA#10525)
Browse files Browse the repository at this point in the history 
* Introducing MegatronStep to make things more extensible

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

* Improve megatron callbacks

Signed-off-by: Marc Romeijn <mromeijn@nvidia.com>
Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

* Some small fixes

Signed-off-by: Marc Romeijn <mromeijn@nvidia.com>
Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

* Apply isort and black reformatting

Signed-off-by: akoumpa <akoumpa@users.noreply.github.com>
Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

* remove debg code

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

* add forward_only to forward backward func

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

* add global-batch-sampler support to MegatronStep

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

* Apply isort and black reformatting

Signed-off-by: akoumpa <akoumpa@users.noreply.github.com>

---------

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
Signed-off-by: Marc Romeijn <mromeijn@nvidia.com>
Signed-off-by: akoumpa <akoumpa@users.noreply.github.com>
Co-authored-by: Marc Romeijn <mromeijn@nvidia.com>
Co-authored-by: akoumpa <akoumpa@users.noreply.github.com>
Co-authored-by: sichu <sichu@nvidia.com>
  • Loading branch information
@akoumpa @marcromeyn @sichu2023
4 people authored Sep 25, 2024
1 parent 0ec10d2 commit e8304d6
Show file tree
Hide file tree
Showing 3 changed files with 550 additions and 225 deletions.
Loading

0 comments on commit e8304d6

Please sign in to comment.