[NeVa Pretraining] Vision Encoder Created on All GPUs During Pipeline Parallelism
This issue is stale because it has been open for 30 days with no activity. Remove stale label or comment or this will be closed in 7 days. This issue was closed because it has been inactive for 7 days since being marked as stale.
Dear all,

I'm encountering an issue with the vision encoder when using pipeline parallelism. 
The vision encoder appears to be instantiated on all GPUs rather than allocated to GPU 0, which is the first stage of the pipeline.

I tracked the vision encode in `NeMo/nemo/collections/multimodal/models/multimodal_llm/neva/neva_model.py` like this.

```
def create_vision_encoder_and_processor(self, mm_cfg):
    # Initialize vision encoder and freeze it
    if mm_cfg.vision_encoder.get("from_hf", False):
        from transformers import AutoConfig

        config = AutoConfig.from_pretrained(mm_cfg.vision_encoder.from_pretrained)
        if config.architectures[0] == "CLIPVisionModel" or config.architectures[0] == "CLIPModel":
            vision_encoder = CLIPVisionModel.from_pretrained(
                mm_cfg.vision_encoder.from_pretrained,
                torch_dtype=torch.bfloat16,
            ).cuda()
            vision_encoder = vision_encoder.to(torch.bfloat16)
            print(f"vision_encoder is on device: {next(vision_encoder.parameters()).device}")
```
The log shows that the vision encoder is on every device.
Since I used pipeline parallelism, I think the vision encoder must reside in the first stage of the pipeline.
If I misunderstood or my tracking is wrong, let me know.

Thank you.

### Experimental Setup

I used the `nvcr.io/nvidia/nemo:24.07` docker image on testing.

This is the following script to run with NeMo.
```
! torchrun --nproc_per_node=4 \
    /opt/NeMo/examples/multimodal/multimodal_llm/neva/neva_pretrain.py \
    trainer.precision=bf16 \
    trainer.num_nodes=1 \
    trainer.devices=4 \
    trainer.val_check_interval=50 \
    trainer.limit_val_batches=5 \
    trainer.log_every_n_steps=1 \
    trainer.max_steps=1 \
    model.megatron_amp_O2=True \
    model.micro_batch_size=1 \
    model.global_batch_size=16 \
    model.tensor_model_parallel_size=1 \
    model.pipeline_model_parallel_size=4 \
    model.mcore_gpt=True \
    model.transformer_engine=True \
    model.data.data_path=/data/neva/datasets/LLaVA-Pretrain-LCS-558k/blip_laion_cc_sbu_558k.json \
    model.data.image_folder=/data/neva/datasets/LLaVA-Pretrain-LCS-558k/images \
    model.tokenizer.library=sentencepiece \
    model.tokenizer.model=/data/neva/tokenizers/tokenizer_neva.model \
    model.encoder_seq_length=4096 \
    model.num_layers=32 \
    model.hidden_size=4096 \
    model.ffn_hidden_size=11008 \
    model.num_attention_heads=32 \
    model.normalization=rmsnorm \
    model.do_layer_norm_weight_decay=False \
    model.apply_query_key_layer_scaling=True \
    model.bias=False \
    model.activation=fast-swiglu \
    model.headscale=False \
    model.position_embedding_type=rope \
    model.rotary_percentage=1.0 \
    model.num_query_groups=null \
    model.data.num_workers=0 \
    model.mm_cfg.llm.from_pretrained=/data/neva/checkpoints/checkpoints/llama-2-7b-chat.nemo \
    model.mm_cfg.llm.model_type=v1 \
    model.data.conv_template=v1 \
    model.mm_cfg.vision_encoder.from_pretrained='openai/clip-vit-large-patch14' \
    model.mm_cfg.vision_encoder.from_hf=True \
    model.optim.name="fused_adam" \
    exp_manager.create_checkpoint_callback=True \
    exp_manager.checkpoint_callback_params.save_nemo_on_train_end=False \
    exp_manager.create_wandb_logger=False
```