Add support and recipes for HF models via AutoModelForCausalLM #10962

akoumpa · 2024-10-20T00:08:42Z

What does this PR do ?

Can be invoked with nemo or via the example

nemo llm pretrain --factory 'HfAutoModelForCausalLM(model_name="meta-llama/Llama-3.2-1B", num_gpus_per_node=1)'

the example:

cd examples/llm/sft/
python3 hf.py --model=meta-llama/Llama-3.2-1B --devices=1 --strategy=ddp

Collection: [Note which collection this PR will affect]

Changelog

Add specific line by line info of high level changes in this PR.

Usage

You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you add or update any necessary documentation?
Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
- Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

New Feature
Bugfix
Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

Related to # (issue)

examples/llm/sft/hf.py

nemo/lightning/pytorch/optim/pytorch.py

nemo/collections/llm/recipes/hf_auto_model.py

Signed-off-by: Alexandros Koumparoulis <[email protected]>

Signed-off-by: akoumpa <[email protected]>

github-actions · 2024-10-23T05:10:49Z

[🤖]: Hi @akoumpa 👋,

We wanted to let you know that a CICD pipeline for this PR just finished successfully

So it might be time to merge this PR or get some approvals

I'm just a bot so I'll leave it you what to do next.

//cc @pablo-garay @ko3n1g

* initial hf_lit_module Signed-off-by: Alexandros Koumparoulis <[email protected]> * make sft gpt dataset sanity check optional Signed-off-by: Alexandros Koumparoulis <[email protected]> * HF sft example Signed-off-by: Alexandros Koumparoulis <[email protected]> * Rename HfLitModule to HfAutoModel Signed-off-by: Alexandros Koumparoulis <[email protected]> * update default model id Signed-off-by: Alexandros Koumparoulis <[email protected]> * move rank&world_size as params Signed-off-by: Alexandros Koumparoulis <[email protected]> * fix mbs in example Signed-off-by: Alexandros Koumparoulis <[email protected]> * fix for fsdp and logger Signed-off-by: Alexandros Koumparoulis <[email protected]> * make loss_fn configurable Signed-off-by: Alexandros Koumparoulis <[email protected]> * Apply isort and black reformatting Signed-off-by: akoumpa <[email protected]> * remove optim from HfAutoModel Signed-off-by: Alexandros Koumparoulis <[email protected]> * add pytorch native optim Signed-off-by: Alexandros Koumparoulis <[email protected]> * add hfAutoModel pretrain nemorun recipe Signed-off-by: Alexandros Koumparoulis <[email protected]> * remove debug Signed-off-by: Alexandros Koumparoulis <[email protected]> * remove stale imports Signed-off-by: Alexandros Koumparoulis <[email protected]> * remove stale import Signed-off-by: Alexandros Koumparoulis <[email protected]> * rm stale imports Signed-off-by: Alexandros Koumparoulis <[email protected]> * rm stale imports Signed-off-by: Alexandros Koumparoulis <[email protected]> * tokenizer fix Signed-off-by: Alexandros Koumparoulis <[email protected]> * update example Signed-off-by: Alexandros Koumparoulis <[email protected]> * rename pytorch_adam to pytorch_adam_with_cosine_annealing Signed-off-by: Alexandros Koumparoulis <[email protected]> * small refactor Signed-off-by: Alexandros Koumparoulis <[email protected]> * fix no_weight_decay_cond Signed-off-by: Alexandros Koumparoulis <[email protected]> * fix Signed-off-by: Alexandros Koumparoulis <[email protected]> * fix Signed-off-by: Alexandros Koumparoulis <[email protected]> * switch to flat_lr optim for example Signed-off-by: Alexandros Koumparoulis <[email protected]> * Apply isort and black reformatting Signed-off-by: akoumpa <[email protected]> * remove imports & update docstrings Signed-off-by: Alexandros Koumparoulis <[email protected]> * add a tokenizer setter to allow it to work with nemo/collections/llm/api.py::_use_tokenizer Signed-off-by: Alexandros Koumparoulis <[email protected]> * remove unused import Signed-off-by: Alexandros Koumparoulis <[email protected]> * allow loss_mask to be none Signed-off-by: Alexandros Koumparoulis <[email protected]> * Add HF-dataset lightning module Signed-off-by: Alexandros Koumparoulis <[email protected]> * check if pad_token_id is None Signed-off-by: Alexandros Koumparoulis <[email protected]> * rename hf_lit_module.py to hf_auto_model.py Signed-off-by: Alexandros Koumparoulis <[email protected]> * class rename Signed-off-by: Alexandros Koumparoulis <[email protected]> * rename Signed-off-by: Alexandros Koumparoulis <[email protected]> * update example Signed-off-by: Alexandros Koumparoulis <[email protected]> * HfAutoModelForCausalLM Signed-off-by: Alexandros Koumparoulis <[email protected]> * rm stale import Signed-off-by: Alexandros Koumparoulis <[email protected]> * add option to start with random weights Signed-off-by: Alexandros Koumparoulis <[email protected]> * add check in megatron-strategy Signed-off-by: Alexandros Koumparoulis <[email protected]> * rename param Signed-off-by: Alexandros Koumparoulis <[email protected]> * drop mcore sampler from squadmodule Signed-off-by: Alexandros Koumparoulis <[email protected]> * make megatron_sampler optional in HfDatasetDataModule Signed-off-by: Alexandros Koumparoulis <[email protected]> * copyright Signed-off-by: Alexandros Koumparoulis <[email protected]> * use is_hf_model to mark hf classes Signed-off-by: Alexandros Koumparoulis <[email protected]> * Apply isort and black reformatting Signed-off-by: akoumpa <[email protected]> --------- Signed-off-by: Alexandros Koumparoulis <[email protected]> Signed-off-by: akoumpa <[email protected]> Co-authored-by: akoumpa <[email protected]> Signed-off-by: Alexandros Koumparoulis <[email protected]>

* initial hf_lit_module Signed-off-by: Alexandros Koumparoulis <[email protected]> * make sft gpt dataset sanity check optional Signed-off-by: Alexandros Koumparoulis <[email protected]> * HF sft example Signed-off-by: Alexandros Koumparoulis <[email protected]> * Rename HfLitModule to HfAutoModel Signed-off-by: Alexandros Koumparoulis <[email protected]> * update default model id Signed-off-by: Alexandros Koumparoulis <[email protected]> * move rank&world_size as params Signed-off-by: Alexandros Koumparoulis <[email protected]> * fix mbs in example Signed-off-by: Alexandros Koumparoulis <[email protected]> * fix for fsdp and logger Signed-off-by: Alexandros Koumparoulis <[email protected]> * make loss_fn configurable Signed-off-by: Alexandros Koumparoulis <[email protected]> * Apply isort and black reformatting Signed-off-by: akoumpa <[email protected]> * remove optim from HfAutoModel Signed-off-by: Alexandros Koumparoulis <[email protected]> * add pytorch native optim Signed-off-by: Alexandros Koumparoulis <[email protected]> * add hfAutoModel pretrain nemorun recipe Signed-off-by: Alexandros Koumparoulis <[email protected]> * remove debug Signed-off-by: Alexandros Koumparoulis <[email protected]> * remove stale imports Signed-off-by: Alexandros Koumparoulis <[email protected]> * remove stale import Signed-off-by: Alexandros Koumparoulis <[email protected]> * rm stale imports Signed-off-by: Alexandros Koumparoulis <[email protected]> * rm stale imports Signed-off-by: Alexandros Koumparoulis <[email protected]> * tokenizer fix Signed-off-by: Alexandros Koumparoulis <[email protected]> * update example Signed-off-by: Alexandros Koumparoulis <[email protected]> * rename pytorch_adam to pytorch_adam_with_cosine_annealing Signed-off-by: Alexandros Koumparoulis <[email protected]> * small refactor Signed-off-by: Alexandros Koumparoulis <[email protected]> * fix no_weight_decay_cond Signed-off-by: Alexandros Koumparoulis <[email protected]> * fix Signed-off-by: Alexandros Koumparoulis <[email protected]> * fix Signed-off-by: Alexandros Koumparoulis <[email protected]> * switch to flat_lr optim for example Signed-off-by: Alexandros Koumparoulis <[email protected]> * Apply isort and black reformatting Signed-off-by: akoumpa <[email protected]> * remove imports & update docstrings Signed-off-by: Alexandros Koumparoulis <[email protected]> * add a tokenizer setter to allow it to work with nemo/collections/llm/api.py::_use_tokenizer Signed-off-by: Alexandros Koumparoulis <[email protected]> * remove unused import Signed-off-by: Alexandros Koumparoulis <[email protected]> * allow loss_mask to be none Signed-off-by: Alexandros Koumparoulis <[email protected]> * Add HF-dataset lightning module Signed-off-by: Alexandros Koumparoulis <[email protected]> * check if pad_token_id is None Signed-off-by: Alexandros Koumparoulis <[email protected]> * rename hf_lit_module.py to hf_auto_model.py Signed-off-by: Alexandros Koumparoulis <[email protected]> * class rename Signed-off-by: Alexandros Koumparoulis <[email protected]> * rename Signed-off-by: Alexandros Koumparoulis <[email protected]> * update example Signed-off-by: Alexandros Koumparoulis <[email protected]> * HfAutoModelForCausalLM Signed-off-by: Alexandros Koumparoulis <[email protected]> * rm stale import Signed-off-by: Alexandros Koumparoulis <[email protected]> * add option to start with random weights Signed-off-by: Alexandros Koumparoulis <[email protected]> * add check in megatron-strategy Signed-off-by: Alexandros Koumparoulis <[email protected]> * rename param Signed-off-by: Alexandros Koumparoulis <[email protected]> * drop mcore sampler from squadmodule Signed-off-by: Alexandros Koumparoulis <[email protected]> * make megatron_sampler optional in HfDatasetDataModule Signed-off-by: Alexandros Koumparoulis <[email protected]> * copyright Signed-off-by: Alexandros Koumparoulis <[email protected]> * use is_hf_model to mark hf classes Signed-off-by: Alexandros Koumparoulis <[email protected]> * Apply isort and black reformatting Signed-off-by: akoumpa <[email protected]> --------- Signed-off-by: Alexandros Koumparoulis <[email protected]> Signed-off-by: akoumpa <[email protected]> Co-authored-by: akoumpa <[email protected]>

…A#10962) * initial hf_lit_module Signed-off-by: Alexandros Koumparoulis <[email protected]> * make sft gpt dataset sanity check optional Signed-off-by: Alexandros Koumparoulis <[email protected]> * HF sft example Signed-off-by: Alexandros Koumparoulis <[email protected]> * Rename HfLitModule to HfAutoModel Signed-off-by: Alexandros Koumparoulis <[email protected]> * update default model id Signed-off-by: Alexandros Koumparoulis <[email protected]> * move rank&world_size as params Signed-off-by: Alexandros Koumparoulis <[email protected]> * fix mbs in example Signed-off-by: Alexandros Koumparoulis <[email protected]> * fix for fsdp and logger Signed-off-by: Alexandros Koumparoulis <[email protected]> * make loss_fn configurable Signed-off-by: Alexandros Koumparoulis <[email protected]> * Apply isort and black reformatting Signed-off-by: akoumpa <[email protected]> * remove optim from HfAutoModel Signed-off-by: Alexandros Koumparoulis <[email protected]> * add pytorch native optim Signed-off-by: Alexandros Koumparoulis <[email protected]> * add hfAutoModel pretrain nemorun recipe Signed-off-by: Alexandros Koumparoulis <[email protected]> * remove debug Signed-off-by: Alexandros Koumparoulis <[email protected]> * remove stale imports Signed-off-by: Alexandros Koumparoulis <[email protected]> * remove stale import Signed-off-by: Alexandros Koumparoulis <[email protected]> * rm stale imports Signed-off-by: Alexandros Koumparoulis <[email protected]> * rm stale imports Signed-off-by: Alexandros Koumparoulis <[email protected]> * tokenizer fix Signed-off-by: Alexandros Koumparoulis <[email protected]> * update example Signed-off-by: Alexandros Koumparoulis <[email protected]> * rename pytorch_adam to pytorch_adam_with_cosine_annealing Signed-off-by: Alexandros Koumparoulis <[email protected]> * small refactor Signed-off-by: Alexandros Koumparoulis <[email protected]> * fix no_weight_decay_cond Signed-off-by: Alexandros Koumparoulis <[email protected]> * fix Signed-off-by: Alexandros Koumparoulis <[email protected]> * fix Signed-off-by: Alexandros Koumparoulis <[email protected]> * switch to flat_lr optim for example Signed-off-by: Alexandros Koumparoulis <[email protected]> * Apply isort and black reformatting Signed-off-by: akoumpa <[email protected]> * remove imports & update docstrings Signed-off-by: Alexandros Koumparoulis <[email protected]> * add a tokenizer setter to allow it to work with nemo/collections/llm/api.py::_use_tokenizer Signed-off-by: Alexandros Koumparoulis <[email protected]> * remove unused import Signed-off-by: Alexandros Koumparoulis <[email protected]> * allow loss_mask to be none Signed-off-by: Alexandros Koumparoulis <[email protected]> * Add HF-dataset lightning module Signed-off-by: Alexandros Koumparoulis <[email protected]> * check if pad_token_id is None Signed-off-by: Alexandros Koumparoulis <[email protected]> * rename hf_lit_module.py to hf_auto_model.py Signed-off-by: Alexandros Koumparoulis <[email protected]> * class rename Signed-off-by: Alexandros Koumparoulis <[email protected]> * rename Signed-off-by: Alexandros Koumparoulis <[email protected]> * update example Signed-off-by: Alexandros Koumparoulis <[email protected]> * HfAutoModelForCausalLM Signed-off-by: Alexandros Koumparoulis <[email protected]> * rm stale import Signed-off-by: Alexandros Koumparoulis <[email protected]> * add option to start with random weights Signed-off-by: Alexandros Koumparoulis <[email protected]> * add check in megatron-strategy Signed-off-by: Alexandros Koumparoulis <[email protected]> * rename param Signed-off-by: Alexandros Koumparoulis <[email protected]> * drop mcore sampler from squadmodule Signed-off-by: Alexandros Koumparoulis <[email protected]> * make megatron_sampler optional in HfDatasetDataModule Signed-off-by: Alexandros Koumparoulis <[email protected]> * copyright Signed-off-by: Alexandros Koumparoulis <[email protected]> * use is_hf_model to mark hf classes Signed-off-by: Alexandros Koumparoulis <[email protected]> * Apply isort and black reformatting Signed-off-by: akoumpa <[email protected]> --------- Signed-off-by: Alexandros Koumparoulis <[email protected]> Signed-off-by: akoumpa <[email protected]> Co-authored-by: akoumpa <[email protected]> Signed-off-by: Hainan Xu <[email protected]>

…A#10962) * initial hf_lit_module Signed-off-by: Alexandros Koumparoulis <[email protected]> * make sft gpt dataset sanity check optional Signed-off-by: Alexandros Koumparoulis <[email protected]> * HF sft example Signed-off-by: Alexandros Koumparoulis <[email protected]> * Rename HfLitModule to HfAutoModel Signed-off-by: Alexandros Koumparoulis <[email protected]> * update default model id Signed-off-by: Alexandros Koumparoulis <[email protected]> * move rank&world_size as params Signed-off-by: Alexandros Koumparoulis <[email protected]> * fix mbs in example Signed-off-by: Alexandros Koumparoulis <[email protected]> * fix for fsdp and logger Signed-off-by: Alexandros Koumparoulis <[email protected]> * make loss_fn configurable Signed-off-by: Alexandros Koumparoulis <[email protected]> * Apply isort and black reformatting Signed-off-by: akoumpa <[email protected]> * remove optim from HfAutoModel Signed-off-by: Alexandros Koumparoulis <[email protected]> * add pytorch native optim Signed-off-by: Alexandros Koumparoulis <[email protected]> * add hfAutoModel pretrain nemorun recipe Signed-off-by: Alexandros Koumparoulis <[email protected]> * remove debug Signed-off-by: Alexandros Koumparoulis <[email protected]> * remove stale imports Signed-off-by: Alexandros Koumparoulis <[email protected]> * remove stale import Signed-off-by: Alexandros Koumparoulis <[email protected]> * rm stale imports Signed-off-by: Alexandros Koumparoulis <[email protected]> * rm stale imports Signed-off-by: Alexandros Koumparoulis <[email protected]> * tokenizer fix Signed-off-by: Alexandros Koumparoulis <[email protected]> * update example Signed-off-by: Alexandros Koumparoulis <[email protected]> * rename pytorch_adam to pytorch_adam_with_cosine_annealing Signed-off-by: Alexandros Koumparoulis <[email protected]> * small refactor Signed-off-by: Alexandros Koumparoulis <[email protected]> * fix no_weight_decay_cond Signed-off-by: Alexandros Koumparoulis <[email protected]> * fix Signed-off-by: Alexandros Koumparoulis <[email protected]> * fix Signed-off-by: Alexandros Koumparoulis <[email protected]> * switch to flat_lr optim for example Signed-off-by: Alexandros Koumparoulis <[email protected]> * Apply isort and black reformatting Signed-off-by: akoumpa <[email protected]> * remove imports & update docstrings Signed-off-by: Alexandros Koumparoulis <[email protected]> * add a tokenizer setter to allow it to work with nemo/collections/llm/api.py::_use_tokenizer Signed-off-by: Alexandros Koumparoulis <[email protected]> * remove unused import Signed-off-by: Alexandros Koumparoulis <[email protected]> * allow loss_mask to be none Signed-off-by: Alexandros Koumparoulis <[email protected]> * Add HF-dataset lightning module Signed-off-by: Alexandros Koumparoulis <[email protected]> * check if pad_token_id is None Signed-off-by: Alexandros Koumparoulis <[email protected]> * rename hf_lit_module.py to hf_auto_model.py Signed-off-by: Alexandros Koumparoulis <[email protected]> * class rename Signed-off-by: Alexandros Koumparoulis <[email protected]> * rename Signed-off-by: Alexandros Koumparoulis <[email protected]> * update example Signed-off-by: Alexandros Koumparoulis <[email protected]> * HfAutoModelForCausalLM Signed-off-by: Alexandros Koumparoulis <[email protected]> * rm stale import Signed-off-by: Alexandros Koumparoulis <[email protected]> * add option to start with random weights Signed-off-by: Alexandros Koumparoulis <[email protected]> * add check in megatron-strategy Signed-off-by: Alexandros Koumparoulis <[email protected]> * rename param Signed-off-by: Alexandros Koumparoulis <[email protected]> * drop mcore sampler from squadmodule Signed-off-by: Alexandros Koumparoulis <[email protected]> * make megatron_sampler optional in HfDatasetDataModule Signed-off-by: Alexandros Koumparoulis <[email protected]> * copyright Signed-off-by: Alexandros Koumparoulis <[email protected]> * use is_hf_model to mark hf classes Signed-off-by: Alexandros Koumparoulis <[email protected]> * Apply isort and black reformatting Signed-off-by: akoumpa <[email protected]> --------- Signed-off-by: Alexandros Koumparoulis <[email protected]> Signed-off-by: akoumpa <[email protected]> Co-authored-by: akoumpa <[email protected]>

github-actions bot added the NLP label Oct 20, 2024

akoumpa force-pushed the akoumparouli/hf_lit_module branch 2 times, most recently from 5a01f6c to 887ed5a Compare October 21, 2024 21:07

github-advanced-security bot found potential problems Oct 21, 2024

View reviewed changes

akoumpa force-pushed the akoumparouli/hf_lit_module branch 14 times, most recently from 723f74c to 2276451 Compare October 21, 2024 22:47

akoumpa requested review from marcromeyn and hemildesai October 21, 2024 22:49

akoumpa added the Run CICD label Oct 21, 2024

github-advanced-security bot found potential problems Oct 21, 2024

View reviewed changes

nemo/collections/llm/recipes/hf_auto_model.py Fixed Show fixed Hide fixed

akoumpa force-pushed the akoumparouli/hf_lit_module branch 7 times, most recently from f431889 to 22931e2 Compare October 22, 2024 05:57

akoumpa removed the Run CICD label Oct 22, 2024

akoumpa added 12 commits October 22, 2024 14:07

rename hf_lit_module.py to hf_auto_model.py

48cc34e

Signed-off-by: Alexandros Koumparoulis <[email protected]>

class rename

0f85f13

Signed-off-by: Alexandros Koumparoulis <[email protected]>

rename

758a7e9

Signed-off-by: Alexandros Koumparoulis <[email protected]>

update example

7c2689c

Signed-off-by: Alexandros Koumparoulis <[email protected]>

HfAutoModelForCausalLM

dc013a9

Signed-off-by: Alexandros Koumparoulis <[email protected]>

rm stale import

a0cc136

Signed-off-by: Alexandros Koumparoulis <[email protected]>

add option to start with random weights

a6a15db

Signed-off-by: Alexandros Koumparoulis <[email protected]>

add check in megatron-strategy

c815714

Signed-off-by: Alexandros Koumparoulis <[email protected]>

rename param

a8b878c

Signed-off-by: Alexandros Koumparoulis <[email protected]>

drop mcore sampler from squadmodule

9da801b

Signed-off-by: Alexandros Koumparoulis <[email protected]>

make megatron_sampler optional in HfDatasetDataModule

3f8dadc

Signed-off-by: Alexandros Koumparoulis <[email protected]>

copyright

8d1d6aa

Signed-off-by: Alexandros Koumparoulis <[email protected]>

akoumpa force-pushed the akoumparouli/hf_lit_module branch from 2733fda to 8d1d6aa Compare October 22, 2024 21:07

use is_hf_model to mark hf classes

bf8ee45

Signed-off-by: Alexandros Koumparoulis <[email protected]>

akoumpa force-pushed the akoumparouli/hf_lit_module branch from c040e20 to bf8ee45 Compare October 22, 2024 21:22

Apply isort and black reformatting

9520cd1

Signed-off-by: akoumpa <[email protected]>

hemildesai approved these changes Oct 22, 2024

View reviewed changes

akoumpa added Run CICD and removed Run CICD labels Oct 22, 2024

akoumpa merged commit 8f26236 into main Oct 23, 2024
153 of 156 checks passed

akoumpa deleted the akoumparouli/hf_lit_module branch October 23, 2024 05:11

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add support and recipes for HF models via AutoModelForCausalLM #10962

Add support and recipes for HF models via AutoModelForCausalLM #10962

akoumpa commented Oct 20, 2024 •

edited

Loading

github-actions bot commented Oct 23, 2024

Add support and recipes for HF models via AutoModelForCausalLM #10962

Add support and recipes for HF models via AutoModelForCausalLM #10962

Conversation

akoumpa commented Oct 20, 2024 • edited Loading

What does this PR do ?

Changelog

Usage

GitHub Actions CI

Before your PR is "Ready for review"

Who can review?

Additional Information

github-actions bot commented Oct 23, 2024

akoumpa commented Oct 20, 2024 •

edited

Loading