Romeyn/sampler (#10525)

* 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>