Wrap all fp8 extra states in LocalNonpersistentObject #9422
Conversation
Signed-off-by: Jan Baczek <jbaczek@nvidia.com>
jbaczek
force-pushed
the
jbaczek/llm/fix_war_for_fp8_load
branch
from
7f822fa
to
549c541
Compare
|
This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days. |
| @@ -1850,7 +1850,7 @@ def sharded_state_dict(self, prefix: str = '') -> Dict[str, Any]: | |||
|
|
|||
| # WAR: This is a temporary fix to skip loading FP8 parameters for Dot Product Attention | |||
| def skip_fp8_load(x): | |||
| if isinstance(x, ShardedObject) and 'fused_attention' in x.key and '_extra_state' in x.key: | |||
There was a problem hiding this comment.
I'm concerned that this PR basically makes the GPT models ignore the FP8 state for all layers in the checkpoint.
In the meantime I prepared a thorough solution, with an almost merged MCore branch and a corresponding NeMo branch.
@jbaczek could you check (in theory or in practice) if this would solve your problem?
The required flag to set would be model.dist_ckpt_load_sctrictness=log_all
There was a problem hiding this comment.
I think that non-strict loading would solve the problem. I see that this branch is already merged to mcore. When should we expect the sync of repositories, so I could use the public implementation?
There was a problem hiding this comment.
I asked for a public sync, ideally should be available later today, but we don't have an official ETA
|
This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days. |
|
This PR was closed because it has been inactive for 7 days since being marked as stale. |
What does this PR do ?
This PR generalizes FP8 extra state wrapping for all tensors.
Collection: nlp
Before your PR is "Ready for review"
Pre checks:
PR Type: