What does this PR do ?

When using overlap param gather, Mcore does a single param all-gather after each backward pass (before the first forward step). This is problematic when the first forward pass following a backward pass is using the reference model parameters, because we end up doing an all-gather on the reference parameters rather than the model being trained. The solution is to disable overlap param gather when running reference model forward and re-enable after.

Issues

Closes #552

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this 

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...