Skip to content

Update AVLM recipes for NeMo-CI runs #14397

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Aug 5, 2025
Merged

Update AVLM recipes for NeMo-CI runs #14397

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
85d61a6
Update avlm_8b.py
huvunvidia Aug 3, 2025
4dec107
Update avlm_8b.py
huvunvidia Aug 4, 2025
7d016c9
Update avlm_8b.py
huvunvidia Aug 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 20 additions & 7 deletions nemo/collections/avlm/recipes/avlm_8b.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,8 @@
from nemo.collections.avlm import AVLMMockDataModule
from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer
from nemo.collections.llm.peft import LoRA
from nemo.collections.llm.recipes.log.default import tensorboard_logger
from nemo.collections.llm.recipes.finetune_default import nemo_resume
from nemo.collections.llm.recipes.log.default import default_resume, tensorboard_logger
from nemo.collections.llm.recipes.optim.adam import distributed_fused_adam_with_cosine_annealing
from nemo.collections.llm.recipes.precision.mixed_precision import bf16_mixed
from nemo.utils.exp_manager import TimingCallback
Expand Down Expand Up @@ -69,7 +70,13 @@ def finetune_recipe(
num_nodes: int = 1,
num_gpus_per_node: int = 8,
peft_scheme: Optional[str] = 'none',
freeze_modules: Optional[dict] = None,
freeze_modules: Optional[dict] = {
"freeze_language_model": False,
"freeze_vision_model": True,
"freeze_audio_model": True,
"freeze_vision_projection": False,
"freeze_audio_projection": False,
},
) -> run.Partial:
"""
Create a fine-tuning recipe for AVLM 8B model.
Expand Down Expand Up @@ -102,6 +109,7 @@ def finetune_recipe(
pipeline_model_parallel_size=1,
encoder_pipeline_model_parallel_size=0,
pipeline_dtype=torch.bfloat16,
ckpt_async_save=False,
)

trainer = run.Config(
Expand Down Expand Up @@ -143,10 +151,7 @@ def finetune_recipe(
),
log=llm.default_log(dir=dir, name=name, tensorboard_logger=tensorboard_logger(name=name)),
optim=distributed_fused_adam_with_cosine_annealing(max_lr=2.0e-05, min_lr=2.0e-07, warmup_steps=150),
resume=run.Config(
nl.AutoResume,
restore_config=run.Config(nl.RestoreConfig, path=checkpoint_path),
),
resume=nemo_resume(checkpoint_path),
)

if peft_scheme is None or peft_scheme.lower() == 'none':
Expand Down Expand Up @@ -178,7 +183,13 @@ def pretrain_recipe(
num_nodes: int = 1,
num_gpus_per_node: int = 8,
language_model_from_pretrained: Optional[str] = None,
freeze_modules: Optional[dict] = None,
freeze_modules: Optional[dict] = {
"freeze_language_model": True,
"freeze_vision_model": True,
"freeze_audio_model": True,
"freeze_vision_projection": False,
"freeze_audio_projection": False,
},
) -> run.Partial:
"""
Create a Pre-training recipe for AVLM 8B model.
Expand Down Expand Up @@ -210,6 +221,7 @@ def pretrain_recipe(
pipeline_model_parallel_size=1,
encoder_pipeline_model_parallel_size=0,
pipeline_dtype=torch.bfloat16,
ckpt_async_save=False,
)

trainer = run.Config(
Expand Down Expand Up @@ -252,6 +264,7 @@ def pretrain_recipe(
),
log=llm.default_log(dir=dir, name=name, tensorboard_logger=tensorboard_logger(name=name)),
optim=distributed_fused_adam_with_cosine_annealing(max_lr=0.001, min_lr=2.0e-05, warmup_steps=150),
resume=default_resume(),
)

return recipe
Loading