Support NeMo NeVA Model

[athitten]

🚀 Feature

NeMo's NeVa (LLaVa) is a multimodal language model

Initial examine:
Found 49 distinct operations, of which 39 (79.6%) are supported

Work items

• Handling inplace through SSA #145 (but looks like Functionalize in-place ops #584 will be enough for this model).
• Implement _VariableFunctionsClass.empty of torch #331
• Implement torch.device  #327
• Implement _VariableFunctionsClass.max of torch #338
• Impement TensorBase.view_as #328
• Implement _VariableFunctionsClass.baddbmm of torch #326
• Implement TensorBase.is_complex #339
• Implement TensorBase.div_ #340
• Implement TensorBase.__setitem__ #341
• Implement TensorBase.clone #342
• Implement TensorBase.is_cuda #329
• unhashable type: 'TensorProxy' error on NeVA model #601
• TypeError: unwrap_if_dead() incompatible args in NeVA model #660
• maximum recursion depth in NeVA #674
• Dtype mismatch in linear layer #678
• NeVa - ValueError: <object> had an unexpected type <class 'object'> #717
• Support torch.where(condition) with thunder.jit #124
• NeVa::forward - remove device syncs (torch.where) and vectorize over batch dimensions NVIDIA/NeMo#9689
• Indexing with slice and list #643
• tentative: type inference: mismatched dtype in cat operator #750
• Support the scatter operator #752
• Indexing with slice and list #643
• Dtype mismatch in cat: bfloat16 and float16 #812
• enter/exit_autocast of torch.amp.autocast_mode #824
• _empty_transform() got multiple values for argument 'device' #825
• Expected the batch dimensions of a (((384,),)) and the batch dimensions of b (()) to be the same #826
• TypeError: expected 2 arguments, got 1 in megatron transformer #858
• Interpreter error in pure-thunder NeVA path: non-WrappedValue in args #891
• TypeError: rms_forward_affine_mixed_dtypes(): incompatible function arguments #896
• Matmul lowering expands tensor input leading to 100x+ slowdown #872
• Unknown attribute _base inside Megatron core #753
• NeVA bug regarding flash_attn_with_kvcache #1004
• KeyError: 'type' from torch.compile executor #1040
• [NeVa] thunder.core.interpreter.InterpreterError: Encountered exception IndexError: list index out of range while tracing GraphModule #1187
• torch.Generator, in particular set_state #1242
• Constant modeling for getitem #1250
• Don't use nvFuser for 1-op fusions of metadata manipulation ops #1251
• [NeVa] [rank0]: TypeError: matmul(): argument 'input' (position 1) must be Tensor, not TensorProxy #1248
• Strength reduction: fold transpose into a subsequent GEMM call #1252

Running the model

Required data

First download the freely available data and place it in a data directory.

NeMo installation

Dependencies

python3 -m pip install --no-deps \
  huggingface-hub==0.23.2

NeMo branch

To keep the whole thunder team on the same NeMo revisions, and to prevent having a bunch of "modify file to call thunder.jit()" instructions, we temporarily maintain our own branch for thunder. You can grab it by cloning https://github.com/tfogal/NeMo.git. Make sure you have checked out the tfogal/thunder-nemo branch.

To install NeMo, run python3 -m pip install -e . from the root of the checked-out directory.

Running the network

rm -fr foo-neva-train; mkdir -p foo-neva-train
HYDRA_FULL_ERROR=1 \
THUNDER_ANNOTATE_TRACES=1 \
NEMO_THUNDER_NEVA=thunder \
python3 ./examples/multimodal/multimodal_llm/neva/neva_pretrain.py \
    trainer.precision=bf16-mixed \
    model.megatron_amp_O2=True \
    model.mcore_gpt=False \
    trainer.num_nodes=1 \
    trainer.devices=1 \
    trainer.val_check_interval=10 \
    trainer.limit_val_batches=5 \
    trainer.log_every_n_steps=1 \
    ++exp_manager.max_time_per_run=00:00:03:00 \
    trainer.max_steps=20 \
    model.micro_batch_size=2 \
    model.global_batch_size=4 \
    model.tensor_model_parallel_size=1 \
    model.pipeline_model_parallel_size=1 \
    exp_manager.create_checkpoint_callback=False \
    model.data.data_path=./data/multimodal/tiny-neva/dummy.json \
    model.data.image_folder=./data/multimodal/tiny-neva/images \
    model.tokenizer.library=sentencepiece \
    model.tokenizer.model=./data/multimodal/tiny-neva/tokenizer_add_special.model \
    model.num_layers=2 \
    model.hidden_size=5120 \
    model.ffn_hidden_size=13824 \
    model.num_attention_heads=40 \
    model.normalization=rmsnorm \
    model.data.num_workers=0 \
    model.data.conv_template=llama_2 \
    model.mm_cfg.vision_encoder.from_pretrained=openai/clip-vit-large-patch14 \
    model.mm_cfg.llm.from_pretrained=null \
    model.use_flash_attention=false \
    exp_manager.exp_dir=./foo-neva-train

Note that the latest version of the tfogal/thunder-nemo branch allows running with dynamo+thunder by setting NEMO_THUNDER_NEVA=dynamo.

cc @apaz-cli @tfogal

[IvanYashchuk]

Can you share the script for the examine call?

[tfogal]

Can you share the script for the examine call?

@athitten when you have a minute

[athitten]

Adding the updated command to use megatron_amp_O2=True and model.mcore_gpt = True (NeMo models will be defaulting to using models from Megatron, hence this setting). With megatron_amp_O2=True, having precision=bf16 should do mixed precision training with main copy of weights in FP32, but just to be safe also specifying precision=bf16-mixed.

python3 ./examples/multimodal/multimodal_llm/neva/neva_pretrain.py trainer.precision=bf16-mixed model.megatron_amp_O2=True model.mcore_gpt=True  trainer.num_nodes=1 trainer.devices=1 trainer.val_check_interval=10 trainer.limit_val_batches=5 trainer.log_every_n_steps=1 ++exp_manager.max_time_per_run=00:00:03:00 trainer.max_steps=20 model.micro_batch_size=2 model.global_batch_size=4 model.tensor_model_parallel_size=1 model.pipeline_model_parallel_size=1 exp_manager.create_checkpoint_callback=False model.data.data_path=./data/multimodal/tiny-neva/dummy.json model.data.image_folder=./data/multimodal/tiny-neva/images model.tokenizer.library=sentencepiece model.tokenizer.model=./data/multimodal/tiny-neva/tokenizer_add_special.model model.num_layers=2 model.hidden_size=5120 model.ffn_hidden_size=13824 model.num_attention_heads=40 model.normalization=rmsnorm model.data.num_workers=0 model.data.conv_template=llama_2 model.mm_cfg.vision_encoder.from_pretrained=openai/clip-vit-large-patch14 model.mm_cfg.llm.from_pretrained=null model.use_flash_attention=false exp_manager.exp_dir=./foo-neva-train

[athitten]

This might be helpful: The full config with default values for all parameters can be found: here. Only the parameters we specify in the run command get overwritten by the specified values and others default to values mentioned in the config.

[tfogal]

Adding the updated command

Thanks, @athitten !
I have edited the original issue to mostly reflect the updated command. Unfortunately #753 blocks setting model.mcore_gpt=True, so for now that one's still False... but let's prioritize that one!

[athitten]

Yes its important to prioritize getting thunder working with mcore_gpt=True as it will be default for NeMo models once we deprecate the legacy path.