NeMo1 -> NeMo2 checkpoint conversion

[jstjohn]

Summary

Nemo1 to nemo2 checkpoint conversion

Details

• Bugfixes to NeMo Checkpoint connector bugfixes NeMo#10647
• Test showing that we can convert a model from nemo1 to nemo2 format and that it was properly initialized, on Geneformer.
• TODO same test for ESM2 once TE conversion is complete Migrate ESM2 to transformer engine #199

Usage

py.test sub-packages/bionemo-geneformer/tests/bionemo/geneformer/test_model.py::test_nemo1_checkpoint_conversion

This could either be used in an interactive session or placed in a script to one off convert some specific checkpoint in nemo1 format to a checkpoint in nemo2 format.

Testing

Test checks that the checkpoint can be converted by this function, and then pointing a model at the new nemo2 checkpoint works as expected when doing fine-tuning resumption.

Tests for these changes can be run via:

py.test sub-packages/bionemo-geneformer/tests/bionemo/geneformer/test_model.py::test_nemo1_checkpoint_conversion

• Did you review the Before your PR is "Ready for review" section before asking for review?
• Did you make sure your changes have tests? Did you test your changes locally?
• Can you add the SKIP_CI label to your PR?
• Can you add the PYTEST_NOT_REQUIRED label to your PR?
• Can you add the JET_NOT_REQUIRED label to your PR?

[jstjohn]

/build-ci

[jstjohn]

/build-ci

[jstjohn]

/build-ci

[jstjohn]

Does this work for you? I don’t think the state dict loading line actually replaces the meta tensor. I think contents are filled in. import torchimport torch.nn as nn# Define a module with parameters initialized as meta tensorsclass MyModule(nn.Module):    def __init__(self):        super(MyModule, self).__init__()        self.linear = nn.Linear(10, 10, device='meta')  # Meta device# Instantiate the modulemodule = MyModule()# Print the current parameter device (meta)print(module.linear.weight.device)  # Should output "meta"# Create a state_dict with actual tensorsstate_dict = {    'linear.weight': torch.randn(10, 10),    'linear.bias': torch.randn(10)}# Load the state_dict into the modulemodule.load_state_dict(state_dict)# Check that the parameters have been replaced with actual tensorsprint(module.linear.weight.device)  # Should output "cpu" or the device of the tensors in state_dictThis prints meta device. And there’s a warning (yikes!) not an error! Sent from my iPhoneOn Sep 26, 2024, at 7:27 PM, Alexandros Koumparoulis ***@***.***> wrote: @akoumpa commented on this pull request. In sub-packages/bionemo-llm/src/bionemo/llm/model/biobert/connector.py:

+ def is_te_mapping(model: BioBertLightningModule) -> bool:

+ """Check for TE layers, for now infer this from the config.""" + return model.config.biobert_spec_option in { + BiobertSpecOption.bert_layer_with_transformer_engine_spec, + BiobertSpecOption.bert_layer_with_transformer_engine_and_qk_ln_spec, + } + + def convert_state(self, source: Dict[str, torch.Tensor], target: BioBertLightningModule) -> BioBertLightningModule: + """Convert the input state_dict keys from nemo1 biobert to nemo2 biobert.""" + te_mapping = self.is_te_mapping(target) # check for TE layers. + target.module.cpu() + new_state_dict_from_old = {} + for k, v in source.items(): + new_key = nemo1_to_nemo2_biobert_key_mapping(k, new_model_prefix="", te_mapping=te_mapping) + new_state_dict_from_old[new_key] = v + target.module.load_state_dict(new_state_dict_from_old, strict=not te_mapping) @jstjohn I would add here something like the following meta_tensors = list(filter(lambda x: isinstance(x[1], torch.Tensor) and x[1].device.type == 'meta', target.module.state_dict().items())) assert len(meta_tensors) == 0, meta_tensors This should print all the tensors that have meta device. The assumption here was that the input state_dict (in this case new_state_dict_from_old) contains all the parameters needed by target.module. Please let me know if that works. The other thing would be add a parameter to nemo_setup to init using CPU instead of meta, but I want to avoid this, for large models (e.g. 100B parameters) this takes too long and it's not useful since the initialized parameters will be overwritten with the ones from the checkpoint. —Reply to this email directly, view it on GitHub, or unsubscribe.You are receiving this because you were mentioned.Message ID: ***@***.***>

[jstjohn]

Haha chat gpt wrote that code and was all like “this will work”. I asked it to run it. It printed meta device. Anyways I manually tried running this in a collab instance since I’m AFK, and I can confirm that it just prints a warning and doesn’t actually fill in the tensors. Sent from my iPhoneOn Sep 26, 2024, at 7:47 PM, John St John ***@***.***> wrote:Does this work for you? I don’t think the state dict loading line actually replaces the meta tensor. I think contents are filled in. import torchimport torch.nn as nn# Define a module with parameters initialized as meta tensorsclass MyModule(nn.Module):    def __init__(self):        super(MyModule, self).__init__()        self.linear = nn.Linear(10, 10, device='meta')  # Meta device# Instantiate the modulemodule = MyModule()# Print the current parameter device (meta)print(module.linear.weight.device)  # Should output "meta"# Create a state_dict with actual tensorsstate_dict = {    'linear.weight': torch.randn(10, 10),    'linear.bias': torch.randn(10)}# Load the state_dict into the modulemodule.load_state_dict(state_dict)# Check that the parameters have been replaced with actual tensorsprint(module.linear.weight.device)  # Should output "cpu" or the device of the tensors in state_dictThis prints meta device. And there’s a warning (yikes!) not an error! Sent from my iPhoneOn Sep 26, 2024, at 7:27 PM, Alexandros Koumparoulis ***@***.***> wrote: @akoumpa commented on this pull request. In sub-packages/bionemo-llm/src/bionemo/llm/model/biobert/connector.py:

+ def is_te_mapping(model: BioBertLightningModule) -> bool:

+ """Check for TE layers, for now infer this from the config.""" + return model.config.biobert_spec_option in { + BiobertSpecOption.bert_layer_with_transformer_engine_spec, + BiobertSpecOption.bert_layer_with_transformer_engine_and_qk_ln_spec, + } + + def convert_state(self, source: Dict[str, torch.Tensor], target: BioBertLightningModule) -> BioBertLightningModule: + """Convert the input state_dict keys from nemo1 biobert to nemo2 biobert.""" + te_mapping = self.is_te_mapping(target) # check for TE layers. + target.module.cpu() + new_state_dict_from_old = {} + for k, v in source.items(): + new_key = nemo1_to_nemo2_biobert_key_mapping(k, new_model_prefix="", te_mapping=te_mapping) + new_state_dict_from_old[new_key] = v + target.module.load_state_dict(new_state_dict_from_old, strict=not te_mapping) @jstjohn I would add here something like the following meta_tensors = list(filter(lambda x: isinstance(x[1], torch.Tensor) and x[1].device.type == 'meta', target.module.state_dict().items())) assert len(meta_tensors) == 0, meta_tensors This should print all the tensors that have meta device. The assumption here was that the input state_dict (in this case new_state_dict_from_old) contains all the parameters needed by target.module. Please let me know if that works. The other thing would be add a parameter to nemo_setup to init using CPU instead of meta, but I want to avoid this, for large models (e.g. 100B parameters) this takes too long and it's not useful since the initialized parameters will be overwritten with the ones from the checkpoint. —Reply to this email directly, view it on GitHub, or unsubscribe.You are receiving this because you were mentioned.Message ID: ***@***.***>

[jstjohn]

/build-ci

[jstjohn]

/build-ci

[jstjohn]

Yes it’s newer than TOT. It’s on the commit I’m working on on the Nemo side to fix the checkpoint conversion stuff. Sent from my iPhoneOn Sep 30, 2024, at 7:01 PM, Peter St. John ***@***.***> wrote: @pstjohn commented on this pull request. On 3rdparty/NeMo: Is this on TOT? Let's just not downgrade if the dependabot updates bring us to a more recent commit —Reply to this email directly, view it on GitHub, or unsubscribe.You are receiving this because you were mentioned.Message ID: ***@***.***>

[jstjohn]

/build-ci

[jstjohn]

/build-ci

[jstjohn]

TODO add the steps to create the .tar.gz files:

cd esm2_650M_nemo2 && tar czvf ../esm2_650M_nemo2.tar.gz *

for example.

[jstjohn]

/build-ci