Add llama3 and distributed checkpoint support in NeVA

[yaoyu-33]

What does this PR do ?

Add llama3 and distributed checkpoint support in NeVA

Collection: [multimodal]

Changelog

• Add specific line by line info of high level changes in this PR.

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this 

GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

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 add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

• New Feature
• Bugfix
• Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

• Related to # (issue)

[cuichenx]

then should we get rid of it..?

[cuichenx]

should we rename this to nemotron?

[yaoyu-33]

that will break many old checkpoint, we might just keep this way for a while...

[mikolajblaz]

I assume this is a temporary implementation for a strict flag because we also need to notify the user which keys are skipped.

Also, this will work only with Zarr ckpt format, for PyT Distributed it will be different

[mikolajblaz]

Also, I don't think it's correct, because we should use the sharded key, not the state dict key (also doesn't account for nested dicts).
Can you try something like this (I didn't run this code so there might be errors)?
This should work for all backends, nested dicts and use correct keys.

from megatron.core.dist_checkpointing.dict_utils import extract_matching_values
from megatron.core.dist_checkpointing.mapping import ShardedBase


    if not strict:
        sharded_state_dict = self.adjust_non_strict_load(path, sharded_state_dict)
    ...

    def adjust_non_strict_load(self, path: _PATH, sharded_state_dict: Dict[str, Any]):
        ckpt_sharded_metadata = dist_checkpointing.load_tensors_metadata(path)
        loaded_keys = []
        missing_keys = []
        unexpected_keys = []

        def should_remove_missing_sharded_base(x: Any):
            if isinstance(x, ShardedBase):
                if x.key in ckpt_sharded_metadata:
                    loaded_keys.append(x.key)
                    return False
                else:
                    unexpected_keys.append(x.key)
                    return True
            return False

        _, sharded_state_dict = extract_matching_values(sharded_state_dict, should_remove_missing_sharded_base)
        logging.info(f'The following keys are not in the checkpoint and will not be loaded: {unexpected_keys}')

        # TODO: compute missing_keys by:
        #  1. all_gather_object of loaded_keys
        #  2. missing_keys = ckpt_sharded_metadata.keys() - loaded_keys
        return sharded_state_dict

[mikolajblaz]

Is weight not needed at all?

[yaoyu-33]

weight already be take care of in super

[mikolajblaz]

I think this can be replaced with

from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint
     ...
     sharded_state_dict.update(make_sharded_tensors_for_checkpoint(state_dict))