Add VLM PEFT test running with 4bit model weights.

.github/workflows/cicd-main.yml

@@ -3626,6 +3626,15 @@ jobs:
       AFTER_SCRIPT: |
         rm -rf nemo_experiments
 
+  L2_VLM_HF_Transformer_PEFT_4bit:
+    needs: [ cicd-test-container-setup ]
+    uses: ./.github/workflows/_test_template.yml
+    if: contains(fromJSON(needs.cicd-test-container-setup.outputs.test_to_run), 'L2_VLM_HF_Transformer_PEFT_4bit') || needs.cicd-test-container-setup.outputs.all == 'true'
+    with:
+      RUNNER: self-hosted-azure-gpus-1
+      SCRIPT: |
+        TRANSFORMERS_OFFLINE=1 python tests/collections/vlm/hf/peft.py --model /home/TestData/vlm/qwen2-2b/ --max-steps 3 --disable-ckpt --use-4bit
+
   L2_VLM_HF_Transformer_PEFT_FSDP:
     needs: [ cicd-test-container-setup ]
     uses: ./.github/workflows/_test_template.yml
@@ -3634,8 +3643,6 @@ jobs:
       RUNNER: self-hosted-azure
       SCRIPT: |
         TRANSFORMERS_OFFLINE=1 python tests/collections/vlm/hf/peft.py --model /home/TestData/vlm/qwen2-2b/ --max-steps 3 --disable-ckpt --strategy fsdp --devices 2
-      AFTER_SCRIPT: |
-        rm -rf nemo_experiments
 
   L2_HF_Transformer_PEFT:
     needs: [ cicd-test-container-setup ]
@@ -4917,6 +4924,7 @@ jobs:
       - L2_HF_Transformer_SFT_nemorun
       - L2_HF_Transformer_SFT_2gpu
       - L2_VLM_HF_Transformer_PEFT
+      - L2_VLM_HF_Transformer_PEFT_4bit
       - L2_VLM_HF_Transformer_PEFT_FSDP
       - L2_HF_Transformer_SFT_2gpu_nemorun
       - L2_HF_Transformer_SFT_TE_Acceleration

examples/vlm/hf/peft.py

@@ -85,6 +85,7 @@ def fmt(sample):
     parser.add_argument('--accelerator', default='gpu', choices=['gpu'])
     parser.add_argument('--max-steps', type=int, default=100)
     parser.add_argument('--wandb-project', type=str, default=None)
+    parser.add_argument('--use-4bit', help="Load model in 4bit", action="store_true")
     args = parser.parse_args()
 
     wandb = None
@@ -103,7 +104,7 @@ def fmt(sample):
     processor = vlm.HFAutoModelForImageTextToText.configure_processor(args.model)
 
     llm.api.finetune(
-        model=vlm.HFAutoModelForImageTextToText(args.model),
+        model=vlm.HFAutoModelForImageTextToText(args.model, load_in_4bit=args.use_4bit),
         data=mk_hf_vlm_dataset(processor, args.mbs, args.gbs),
         trainer=nl.Trainer(
             devices=args.devices,
@@ -124,5 +125,6 @@ def fmt(sample):
         peft=llm.peft.LoRA(
             target_modules=['*_proj'],
             dim=16,
+            lora_dtype=torch.bfloat16 if args.use_4bit else None,
         ),
     )

requirements/requirements_multimodal.txt

@@ -1,4 +1,5 @@
 addict
+bitsandbytes==0.45.0
 clip
 decord; sys_platform == 'linux'
 diffusers>=0.19.3

tests/collections/vlm/hf/peft.py

@@ -86,6 +86,7 @@ def fmt(sample):
     parser.add_argument('--max-steps', type=int, default=100)
     parser.add_argument('--wandb-project', type=str, default=None)
     parser.add_argument('--disable-ckpt', action='store_false')
+    parser.add_argument('--use-4bit', help="Load model in 4bit", action="store_true")
     args = parser.parse_args()
 
     wandb = None
@@ -103,7 +104,7 @@ def fmt(sample):
     processor = vlm.HFAutoModelForImageTextToText.configure_processor(args.model)
 
     llm.api.finetune(
-        model=vlm.HFAutoModelForImageTextToText(args.model),
+        model=vlm.HFAutoModelForImageTextToText(args.model, load_in_4bit=args.use_4bit),
         data=mk_hf_vlm_dataset(processor, args.mbs, args.gbs),
         trainer=nl.Trainer(
             devices=args.devices,
@@ -125,5 +126,6 @@ def fmt(sample):
         peft=llm.peft.LoRA(
             target_modules=['*_proj'],
             dim=16,
+            lora_dtype=torch.bfloat16 if args.use_4bit else None,
         ),
     )