Update NeMo/Megatron

[sichu2023]

Pass esm2 golden value tests.

Bugs

• Notice some test leakage, i.e. pass individually but fail when run together
• Fiddle error
• dtype error comes from Megatron-LM

Pytest errors

FAILED scripts/protein/esm2/test_esm2_pretrain.py::test_main_runs - AttributeError: 'NoneType' object has no attribute 'dtype'
FAILED scripts/protein/esm2/test_esm2_pretrain.py::test_val_dataloader_in_main_runs_with_limit_val_batches[1.0] - AttributeError: 'NoneType' object has no attribute 'dtype'
FAILED scripts/protein/esm2/test_esm2_pretrain.py::test_val_dataloader_in_main_runs_with_limit_val_batches[4] - AttributeError: 'NoneType' object has no attribute 'dtype'
FAILED scripts/protein/esm2/test_esm2_pretrain.py::test_val_dataloader_in_main_runs_with_limit_val_batches[None] - AttributeError: 'NoneType' object has no attribute 'dtype'
FAILED scripts/singlecell/geneformer/test_train.py::test_pretrain_cli - AssertionError: Pretrain script failed: python      /workspaces/bionemo-fw-ea/scripts/singlecell/geneformer/train.py         --data-dir /home/bionemo/.cache/bionemo/7bd3714b68c30c50f0a36a02fa9dd720-singlecell-testdata-20240506.tar.gz.untar/cellxgen...
FAILED sub-packages/bionemo-esm2/tests/bionemo/esm2/model/finetune/test_finetune.py::test_esm2_finetune_token_classifier[False] - AttributeError: 'NoneType' object has no attribute 'dtype'
FAILED sub-packages/bionemo-esm2/tests/bionemo/esm2/model/finetune/test_finetune.py::test_esm2_finetune_regressor[False] - AttributeError: 'NoneType' object has no attribute 'dtype'
FAILED sub-packages/bionemo-example_model/tests/bionemo/example_model/test_lightning_basic.py::test_train_mnist_litautoencoder_with_megatron_strategy_single_gpu[32] - AttributeError: 'NoneType' object has no attribute 'dtype'
FAILED sub-packages/bionemo-example_model/tests/bionemo/example_model/test_lightning_basic.py::test_train_mnist_litautoencoder_with_megatron_strategy_single_gpu[bf16-mixed] - AttributeError: 'NoneType' object has no attribute 'dtype'
FAILED sub-packages/bionemo-geneformer/tests/bionemo/geneformer/test_model.py::test_continue_from_checkpoint_geneformer - AttributeError: 'NoneType' object has no attribute 'dtype'
FAILED sub-packages/bionemo-geneformer/tests/bionemo/geneformer/test_model.py::test_finetune_geneformer - AttributeError: 'NoneType' object has no attribute 'dtype'
FAILED sub-packages/bionemo-llm/tests/bionemo/llm/utils/test_iomixin_utils.py::TestIOMixin::test_dataclass_out_of_sync - AssertionError: assert {'b': 7} == {'b': 7, 'c': 3}
FAILED sub-packages/bionemo-llm/tests/bionemo/llm/utils/test_iomixin_utils.py::TestIOMixin::test_dataclass_hparam_modify_parent_default - AssertionError: assert {'a': 7} == {'a': 7, 'b': 3, 'c': 3}
FAILED sub-packages/bionemo-testing/tests/bionemo/testing/data/test_load.py::test_default_pbss_client - botocore.exceptions.ConfigParseError: Unable to parse config file: /home/bionemo/.aws/config

dtype error comes from Megatron-LM

def sharded_param_state_fs_model_space
...
                        for state_key, state_ten in tensors.items():
                            replace_kwargs = dict(
                                key=f'{prefix}.{state_key}.{sharded_metadata.key}',
                                data=state_ten,
>                               dtype=state_ten.dtype,
                                flattened_range=slice(param_range.start, param_range.end),
                                replica_id=replica_id,
                            )
E                           AttributeError: 'NoneType' object has no attribute 'dtype'

3rdparty/Megatron-LM/megatron/core/optimizer/distrib_optimizer.py:1159: AttributeError

[sichu2023]

This is moved outside of if is_te_min_version("1.10.0"). Otherwise we can just call super().__init__ directly without override.

[jstjohn]

This is related in some ways to #304, maybe sync with @farhadrgh and make sure that you are on a new enough NeMo for his needs as well? I think his stuff was recently merged.

[sichu2023]

Dependent on https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/merge_requests/2219