From ee025466ecdbdc8c437e91a5f4b6bd98fd52a190 Mon Sep 17 00:00:00 2001 From: Yash Mehan Date: Tue, 24 Feb 2026 10:38:59 +0000 Subject: [PATCH 01/18] accommodate renaming of AutoModelForVision2Seq to AutoModelForImageTextToText Signed-off-by: Yash Mehan --- tests/utils/test_embedding_resize.py | 5 +++-- tuning/sft_trainer.py | 7 +++++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/utils/test_embedding_resize.py b/tests/utils/test_embedding_resize.py index 5ec7e7ab7c..0c0c7980bf 100644 --- a/tests/utils/test_embedding_resize.py +++ b/tests/utils/test_embedding_resize.py @@ -21,7 +21,7 @@ # Third Party from transformers import ( AutoModelForCausalLM, - AutoModelForVision2Seq, + AutoModelForImageTextToText, #AutoModelForVision2Seq was renamed to this in transformers v5 AutoProcessor, AutoTokenizer, ) @@ -212,7 +212,8 @@ def test_resize_with_multiple_of(): def test_resize_llama_vision_model(): - model = AutoModelForVision2Seq.from_pretrained(TINY_LLAMA_VISION_MODEL_NAME) + # model = AutoModelForVision2Seq.from_pretrained(TINY_LLAMA_VISION_MODEL_NAME) + model = AutoModelForImageTextToText.from_pretrained(TINY_LLAMA_VISION_MODEL_NAME) processor = AutoProcessor.from_pretrained(TINY_LLAMA_VISION_MODEL_NAME) tokenizer = processor.tokenizer diff --git a/tuning/sft_trainer.py b/tuning/sft_trainer.py index 1fa524e896..131b03a6aa 100644 --- a/tuning/sft_trainer.py +++ b/tuning/sft_trainer.py @@ -30,7 +30,7 @@ from torch.cuda import OutOfMemoryError from transformers import ( AutoModelForCausalLM, - AutoModelForVision2Seq, + AutoModelForImageTextToText, # AutoModelForVision2Seq was renamed in transformers v5 AutoProcessor, AutoTokenizer, TrainerCallback, @@ -292,7 +292,10 @@ def train( ) ) # try to load model as a vision model - model = AutoModelForVision2Seq.from_pretrained( + # model = AutoModelForVision2Seq.from_pretrained( + # model_args.model_name_or_path, **model_kwargs + # ) + model = AutoModelForImageTextToText.from_pretrained( model_args.model_name_or_path, **model_kwargs ) try: From 6b797d27905bb529e99458231e653811aaa995d2 Mon Sep 17 00:00:00 2001 From: Yash Mehan Date: Tue, 24 Feb 2026 10:44:42 +0000 Subject: [PATCH 02/18] accomodate renaming of include_tokens_per_second to include_num_input_tokens_seen Signed-off-by: Yash Mehan --- tests/build/test_launch_script.py | 3 ++- tests/test_sft_trainer.py | 6 ++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/build/test_launch_script.py b/tests/build/test_launch_script.py index 322fe5998f..f3bf09dd67 100644 --- a/tests/build/test_launch_script.py +++ b/tests/build/test_launch_script.py @@ -51,7 +51,8 @@ "warmup_ratio": 0.03, "lr_scheduler_type": "cosine", "logging_steps": 1, - "include_tokens_per_second": True, + # "include_tokens_per_second": True, + "include_num_input_tokens_seen": True, "packing": False, "response_template": "\n### Label:", "dataset_text_field": "output", diff --git a/tests/test_sft_trainer.py b/tests/test_sft_trainer.py index c5423fd179..8e6fd50c2c 100644 --- a/tests/test_sft_trainer.py +++ b/tests/test_sft_trainer.py @@ -124,7 +124,8 @@ warmup_ratio=0.03, lr_scheduler_type="cosine", logging_steps=1, - include_tokens_per_second=True, + # include_tokens_per_second=True, + include_num_input_tokens_seen=True, packing=False, max_seq_length=4096, save_strategy="epoch", @@ -140,7 +141,8 @@ warmup_ratio=0.03, lr_scheduler_type="cosine", logging_steps=1, - include_tokens_per_second=True, + # include_tokens_per_second=True, + include_num_input_tokens_seen=True, packing=False, max_seq_length=4096, save_strategy="epoch", From b74972754e75042fdb9b0700d89d5bd14cf51ed1 Mon Sep 17 00:00:00 2001 From: Yash Mehan Date: Tue, 24 Feb 2026 13:15:12 +0000 Subject: [PATCH 03/18] accommodate mllama model.language_model issue and gptneoxtokenizer and llama tokenizer Signed-off-by: Yash Mehan --- tuning/data/tokenizer_utils.py | 43 +++++++++++++++++++++++++++------- 1 file changed, 34 insertions(+), 9 deletions(-) diff --git a/tuning/data/tokenizer_utils.py b/tuning/data/tokenizer_utils.py index 622da0e3c8..303a532395 100644 --- a/tuning/data/tokenizer_utils.py +++ b/tuning/data/tokenizer_utils.py @@ -44,21 +44,42 @@ def get_special_tokens_dict( special_tokens_dict = {} if not tokenizer_name_or_path: - # TODO: understand if we need to hardcode these here or just use defaults in model - if isinstance( - tokenizer, (transformers.LlamaTokenizer, transformers.LlamaTokenizerFast) - ): + # # TODO: understand if we need to hardcode these here or just use defaults in model + # if isinstance( + # tokenizer, (transformers.LlamaTokenizer, transformers.LlamaTokenizerFast) + # ): + llama_classes = tuple( + cls for cls in [ + getattr(transformers, "LlamaTokenizer", None), + getattr(transformers, "LlamaTokenizerFast", None), + ] if cls is not None + ) + is_llama_tokenizer = ( + (bool(llama_classes) and isinstance(tokenizer, llama_classes)) + or "llama" in (getattr(tokenizer, "name_or_path", "") or "").lower() + ) + + gpt_neox_classes = tuple( + cls for cls in [ + getattr(transformers, "GPTNeoXTokenizerFast", None), + getattr(transformers, "GPTNeoXTokenizer", None), + ] if cls is not None + ) + + if is_llama_tokenizer: special_tokens_dict["bos_token"] = "" special_tokens_dict["eos_token"] = "" special_tokens_dict["unk_token"] = "" special_tokens_dict["pad_token"] = "" elif isinstance( - tokenizer, (transformers.GPT2Tokenizer, transformers.GPTNeoXTokenizerFast) + # tokenizer, (transformers.GPT2Tokenizer, transformers.GPTNeoXTokenizerFast) + tokenizer, (transformers.GPT2Tokenizer, *gpt_neox_classes) ): special_tokens_dict["pad_token"] = "" # Add special tokens only when a custom tokenizer is not passed - if tokenizer.pad_token is None: + # if tokenizer.pad_token is None: + if tokenizer.pad_token is None or "pad_token" in special_tokens_dict: logger.warning("PAD token set to default, missing in tokenizer") special_tokens_dict["pad_token"] = configs.DEFAULT_PAD_TOKEN if tokenizer.eos_token is None: @@ -102,7 +123,8 @@ def tokenizer_and_embedding_resize( dict: Metadata on number of added tokens. """ num_new_tokens = tokenizer.add_special_tokens( - special_tokens_dict=special_tokens_dict, replace_additional_special_tokens=False + special_tokens_dict=special_tokens_dict, + # replace_additional_special_tokens=False ) embedding_size = int(multiple_of * math.ceil(len(tokenizer) / multiple_of)) num_new_tokens = num_new_tokens + embedding_size - len(tokenizer) @@ -119,8 +141,11 @@ def tokenizer_and_embedding_resize( model.set_input_embeddings(resized_input_embeddings) # Resize vocab size when embeddings updated for Mllama models - if model.language_model.vocab_size != embedding_size: - model.language_model.vocab_size = embedding_size + # if model.language_model.vocab_size != embedding_size: + # model.language_model.vocab_size = embedding_size + if model.model.vocab_size != embedding_size: + model.model.vocab_size = embedding_size + else: model.resize_token_embeddings(embedding_size) From a3bc6388afd3d582cfaa0f2dc85fb707cbfd5267 Mon Sep 17 00:00:00 2001 From: Yash Mehan Date: Tue, 24 Feb 2026 13:22:14 +0000 Subject: [PATCH 04/18] replace addition_special_tokens with extra_special_tokens Signed-off-by: Yash Mehan --- tests/utils/test_embedding_resize.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/tests/utils/test_embedding_resize.py b/tests/utils/test_embedding_resize.py index 0c0c7980bf..ff2948df90 100644 --- a/tests/utils/test_embedding_resize.py +++ b/tests/utils/test_embedding_resize.py @@ -128,9 +128,10 @@ def test_special_tokens_before_and_after(): model = AutoModelForCausalLM.from_pretrained(MODEL_NAME) input_tokenizer_len = len(tokenizer.get_vocab()) - addn_spl_tokens_before = tokenizer.special_tokens_map.get( - "additional_special_tokens" - ) + # addn_spl_tokens_before = tokenizer.special_tokens_map.get( + # "additional_special_tokens" + # ) + addn_spl_tokens_before = list(tokenizer.extra_special_tokens) assert ( len(addn_spl_tokens_before) > 0 ), "this test needs tokenizer special tokens to not be empty before testing" @@ -150,9 +151,10 @@ def test_special_tokens_before_and_after(): addn_spl_tokens_before.extend(addn_spl_tokens_added) expected_addn_special_tokens = addn_spl_tokens_before expected_embedding_size = input_tokenizer_len + len(addn_spl_tokens_added) + 2 - addn_spl_tokens_after = tokenizer.special_tokens_map.get( - "additional_special_tokens" - ) + # addn_spl_tokens_after = tokenizer.special_tokens_map.get( + # "additional_special_tokens" + # ) + addn_spl_tokens_after = list(tokenizer.extra_special_tokens) assert "" in tokenizer.get_vocab() assert "" in tokenizer.get_vocab() From 28605d4abc7e1e6157852affb048cbb00b1f7f11 Mon Sep 17 00:00:00 2001 From: Yash Mehan Date: Thu, 26 Feb 2026 09:33:52 +0000 Subject: [PATCH 05/18] changed library versions in pyproject.toml Signed-off-by: Yash Mehan --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index a0c5ea7f5b..3f912631e2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,14 +28,14 @@ classifiers=[ dependencies = [ "numpy>=1.26.4,<2.2.0", "accelerate>=1.9.0,<2.0.0", -"transformers>=4.55.0,<=4.55.4", +"transformers==5.2.0", "torch>2.7.0,<2.9.0", -"torchvision<0.24", +"torchvision>=0.25.0", "sentencepiece>=0.1.99,<0.3", "tokenizers<=0.22", "tqdm>=4.66.2,<5.0", "trl>=0.19.1,<0.20.0", -"peft>=0.18.0,< 0.19.0", +"peft>=0.18.1,<0.19.0", "datasets>=4.0.0,<5.0.0", "simpleeval>=0.9.13,<2.0", "pillow>=12.1.1", From ba6ffb0bb31bcdafaa076e0825ec83194b974171 Mon Sep 17 00:00:00 2001 From: Yash Mehan Date: Thu, 26 Feb 2026 11:25:32 +0000 Subject: [PATCH 06/18] fix pyproject toml after the rebase Signed-off-by: Yash Mehan --- pyproject.toml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 3f912631e2..77c53f1fcc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,18 +29,18 @@ dependencies = [ "numpy>=1.26.4,<2.2.0", "accelerate>=1.9.0,<2.0.0", "transformers==5.2.0", -"torch>2.7.0,<2.9.0", +"torch==2.10.0", "torchvision>=0.25.0", "sentencepiece>=0.1.99,<0.3", -"tokenizers<=0.22", +"tokenizers==0.22.2", "tqdm>=4.66.2,<5.0", "trl>=0.19.1,<0.20.0", "peft>=0.18.1,<0.19.0", "datasets>=4.0.0,<5.0.0", "simpleeval>=0.9.13,<2.0", "pillow>=12.1.1", -"kernels<=0.9.0", -"tuning_config_recommender>=0.1.7", +"kernels==0.12.1", +"huggingface_hub>=1.3.0" ] [project.optional-dependencies] From 6c28f16d3ee45cce122f5baa7e66ee84a71ae811 Mon Sep 17 00:00:00 2001 From: Yash Mehan Date: Thu, 26 Feb 2026 11:12:52 +0000 Subject: [PATCH 07/18] add more handing for the filenotfound exception handling in multi rank case Signed-off-by: Yash Mehan --- build/accelerate_launch.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/build/accelerate_launch.py b/build/accelerate_launch.py index 43cf8dda08..5926c4a852 100644 --- a/build/accelerate_launch.py +++ b/build/accelerate_launch.py @@ -116,6 +116,16 @@ def main(): return_code = INTERNAL_ERROR_EXIT_CODE write_termination_log(f"Unhandled exception during training. {e}") sys.exit(return_code) + except Exception as e: # pylint: disable=broad-except + logging.error(traceback.format_exc()) + # v5: torch.distributed raises ChildFailedError with per-rank exit codes + # Check if the root cause was a user error + if hasattr(e, 'failures'): + root_codes = [f.exitcode for f in e.failures.values()] + if any(c == USER_ERROR_EXIT_CODE for c in root_codes): + sys.exit(USER_ERROR_EXIT_CODE) + write_termination_log(f"Unhandled exception during training. {e}") + sys.exit(INTERNAL_ERROR_EXIT_CODE) except Exception as e: # pylint: disable=broad-except logging.error(traceback.format_exc()) write_termination_log(f"Unhandled exception during training. {e}") From 800b2888da434b9d9f2cfe245a643c8000ced52d Mon Sep 17 00:00:00 2001 From: Yash Mehan Date: Thu, 26 Feb 2026 11:13:32 +0000 Subject: [PATCH 08/18] merge existing extra_special_tokens with new ones to prevent replacement Signed-off-by: Yash Mehan --- tests/utils/test_embedding_resize.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/utils/test_embedding_resize.py b/tests/utils/test_embedding_resize.py index ff2948df90..76a68a2040 100644 --- a/tests/utils/test_embedding_resize.py +++ b/tests/utils/test_embedding_resize.py @@ -138,7 +138,10 @@ def test_special_tokens_before_and_after(): special_tokens_dict = {"sep_token": "", "pad_token": ""} addn_spl_tokens_added = ["", "", ""] - special_tokens_dict["additional_special_tokens"] = addn_spl_tokens_added + # special_tokens_dict["additional_special_tokens"] = addn_spl_tokens_added + # for transformers v5: merge existing extra_special_tokens with new ones to prevent replacement + special_tokens_dict["additional_special_tokens"] = list(tokenizer.extra_special_tokens) + addn_spl_tokens_added + resize_result = tokenizer_and_embedding_resize( special_tokens_dict=special_tokens_dict, From c7c9d89ac3296fc837bdc4a498c9593b2e287cb5 Mon Sep 17 00:00:00 2001 From: Yash Mehan Date: Thu, 26 Feb 2026 11:21:01 +0000 Subject: [PATCH 09/18] remove the old commented lines, only letting the new renamed lines stay, adding justification for what was renamed Signed-off-by: Yash Mehan --- tests/utils/test_embedding_resize.py | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/tests/utils/test_embedding_resize.py b/tests/utils/test_embedding_resize.py index 76a68a2040..49b95827dd 100644 --- a/tests/utils/test_embedding_resize.py +++ b/tests/utils/test_embedding_resize.py @@ -128,19 +128,15 @@ def test_special_tokens_before_and_after(): model = AutoModelForCausalLM.from_pretrained(MODEL_NAME) input_tokenizer_len = len(tokenizer.get_vocab()) - # addn_spl_tokens_before = tokenizer.special_tokens_map.get( - # "additional_special_tokens" - # ) - addn_spl_tokens_before = list(tokenizer.extra_special_tokens) + addn_spl_tokens_before = list(tokenizer.extra_special_tokens) # "additional_special_tokens" was renamed to extra_special_tokens in transformers v5 assert ( len(addn_spl_tokens_before) > 0 ), "this test needs tokenizer special tokens to not be empty before testing" special_tokens_dict = {"sep_token": "", "pad_token": ""} addn_spl_tokens_added = ["", "", ""] - # special_tokens_dict["additional_special_tokens"] = addn_spl_tokens_added # for transformers v5: merge existing extra_special_tokens with new ones to prevent replacement - special_tokens_dict["additional_special_tokens"] = list(tokenizer.extra_special_tokens) + addn_spl_tokens_added + special_tokens_dict["additional_special_tokens"] = list(tokenizer.extra_special_tokens) + addn_spl_tokens_added # "additional_special_tokens" was renamed to extra_special_tokens in transformers v5 resize_result = tokenizer_and_embedding_resize( @@ -154,10 +150,7 @@ def test_special_tokens_before_and_after(): addn_spl_tokens_before.extend(addn_spl_tokens_added) expected_addn_special_tokens = addn_spl_tokens_before expected_embedding_size = input_tokenizer_len + len(addn_spl_tokens_added) + 2 - # addn_spl_tokens_after = tokenizer.special_tokens_map.get( - # "additional_special_tokens" - # ) - addn_spl_tokens_after = list(tokenizer.extra_special_tokens) + addn_spl_tokens_after = list(tokenizer.extra_special_tokens) # "additional_special_tokens" was renamed to extra_special_tokens in transformers v5 assert "" in tokenizer.get_vocab() assert "" in tokenizer.get_vocab() @@ -217,8 +210,7 @@ def test_resize_with_multiple_of(): def test_resize_llama_vision_model(): - # model = AutoModelForVision2Seq.from_pretrained(TINY_LLAMA_VISION_MODEL_NAME) - model = AutoModelForImageTextToText.from_pretrained(TINY_LLAMA_VISION_MODEL_NAME) + model = AutoModelForImageTextToText.from_pretrained(TINY_LLAMA_VISION_MODEL_NAME) # AutoModelForVision2Seq was renamed to AutoModelForImageTextToText in transformers v5 processor = AutoProcessor.from_pretrained(TINY_LLAMA_VISION_MODEL_NAME) tokenizer = processor.tokenizer From d5e1862ee30a20a98fe95efc21f97c1723646261 Mon Sep 17 00:00:00 2001 From: Yash Mehan Date: Thu, 26 Feb 2026 11:21:58 +0000 Subject: [PATCH 10/18] remove the old commented lines Signed-off-by: Yash Mehan --- tests/build/test_launch_script.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/build/test_launch_script.py b/tests/build/test_launch_script.py index f3bf09dd67..400ecc105c 100644 --- a/tests/build/test_launch_script.py +++ b/tests/build/test_launch_script.py @@ -51,7 +51,6 @@ "warmup_ratio": 0.03, "lr_scheduler_type": "cosine", "logging_steps": 1, - # "include_tokens_per_second": True, "include_num_input_tokens_seen": True, "packing": False, "response_template": "\n### Label:", From 407a43d43978c4efb531350161ca7fe55353bcb9 Mon Sep 17 00:00:00 2001 From: Yash Mehan Date: Thu, 26 Feb 2026 11:48:32 +0000 Subject: [PATCH 11/18] remove the old commented lines Signed-off-by: Yash Mehan --- tuning/sft_trainer.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tuning/sft_trainer.py b/tuning/sft_trainer.py index 131b03a6aa..51ab291603 100644 --- a/tuning/sft_trainer.py +++ b/tuning/sft_trainer.py @@ -292,9 +292,7 @@ def train( ) ) # try to load model as a vision model - # model = AutoModelForVision2Seq.from_pretrained( - # model_args.model_name_or_path, **model_kwargs - # ) + # in transformers v5, AutoModelForVision2Seq was renamed to AutoModelForImageTextToText model = AutoModelForImageTextToText.from_pretrained( model_args.model_name_or_path, **model_kwargs ) From 68ca44c99ac4137a3447f5f04d6e02d5ca26a0b7 Mon Sep 17 00:00:00 2001 From: Yash Mehan Date: Fri, 27 Feb 2026 05:06:27 +0000 Subject: [PATCH 12/18] complete housekeeping and removing old comments and repetitive comments Signed-off-by: Yash Mehan --- tests/test_sft_trainer.py | 2 -- tests/utils/test_embedding_resize.py | 6 +++--- tuning/data/tokenizer_utils.py | 8 -------- 3 files changed, 3 insertions(+), 13 deletions(-) diff --git a/tests/test_sft_trainer.py b/tests/test_sft_trainer.py index 8e6fd50c2c..df972dfbcd 100644 --- a/tests/test_sft_trainer.py +++ b/tests/test_sft_trainer.py @@ -124,7 +124,6 @@ warmup_ratio=0.03, lr_scheduler_type="cosine", logging_steps=1, - # include_tokens_per_second=True, include_num_input_tokens_seen=True, packing=False, max_seq_length=4096, @@ -141,7 +140,6 @@ warmup_ratio=0.03, lr_scheduler_type="cosine", logging_steps=1, - # include_tokens_per_second=True, include_num_input_tokens_seen=True, packing=False, max_seq_length=4096, diff --git a/tests/utils/test_embedding_resize.py b/tests/utils/test_embedding_resize.py index 49b95827dd..e723b486ba 100644 --- a/tests/utils/test_embedding_resize.py +++ b/tests/utils/test_embedding_resize.py @@ -128,7 +128,7 @@ def test_special_tokens_before_and_after(): model = AutoModelForCausalLM.from_pretrained(MODEL_NAME) input_tokenizer_len = len(tokenizer.get_vocab()) - addn_spl_tokens_before = list(tokenizer.extra_special_tokens) # "additional_special_tokens" was renamed to extra_special_tokens in transformers v5 + addn_spl_tokens_before = list(tokenizer.extra_special_tokens) assert ( len(addn_spl_tokens_before) > 0 ), "this test needs tokenizer special tokens to not be empty before testing" @@ -136,7 +136,7 @@ def test_special_tokens_before_and_after(): special_tokens_dict = {"sep_token": "", "pad_token": ""} addn_spl_tokens_added = ["", "", ""] # for transformers v5: merge existing extra_special_tokens with new ones to prevent replacement - special_tokens_dict["additional_special_tokens"] = list(tokenizer.extra_special_tokens) + addn_spl_tokens_added # "additional_special_tokens" was renamed to extra_special_tokens in transformers v5 + special_tokens_dict["additional_special_tokens"] = list(tokenizer.extra_special_tokens) + addn_spl_tokens_added resize_result = tokenizer_and_embedding_resize( @@ -150,7 +150,7 @@ def test_special_tokens_before_and_after(): addn_spl_tokens_before.extend(addn_spl_tokens_added) expected_addn_special_tokens = addn_spl_tokens_before expected_embedding_size = input_tokenizer_len + len(addn_spl_tokens_added) + 2 - addn_spl_tokens_after = list(tokenizer.extra_special_tokens) # "additional_special_tokens" was renamed to extra_special_tokens in transformers v5 + addn_spl_tokens_after = list(tokenizer.extra_special_tokens) assert "" in tokenizer.get_vocab() assert "" in tokenizer.get_vocab() diff --git a/tuning/data/tokenizer_utils.py b/tuning/data/tokenizer_utils.py index 303a532395..eb5ce5fb8d 100644 --- a/tuning/data/tokenizer_utils.py +++ b/tuning/data/tokenizer_utils.py @@ -44,10 +44,6 @@ def get_special_tokens_dict( special_tokens_dict = {} if not tokenizer_name_or_path: - # # TODO: understand if we need to hardcode these here or just use defaults in model - # if isinstance( - # tokenizer, (transformers.LlamaTokenizer, transformers.LlamaTokenizerFast) - # ): llama_classes = tuple( cls for cls in [ getattr(transformers, "LlamaTokenizer", None), @@ -72,13 +68,11 @@ def get_special_tokens_dict( special_tokens_dict["unk_token"] = "" special_tokens_dict["pad_token"] = "" elif isinstance( - # tokenizer, (transformers.GPT2Tokenizer, transformers.GPTNeoXTokenizerFast) tokenizer, (transformers.GPT2Tokenizer, *gpt_neox_classes) ): special_tokens_dict["pad_token"] = "" # Add special tokens only when a custom tokenizer is not passed - # if tokenizer.pad_token is None: if tokenizer.pad_token is None or "pad_token" in special_tokens_dict: logger.warning("PAD token set to default, missing in tokenizer") special_tokens_dict["pad_token"] = configs.DEFAULT_PAD_TOKEN @@ -141,8 +135,6 @@ def tokenizer_and_embedding_resize( model.set_input_embeddings(resized_input_embeddings) # Resize vocab size when embeddings updated for Mllama models - # if model.language_model.vocab_size != embedding_size: - # model.language_model.vocab_size = embedding_size if model.model.vocab_size != embedding_size: model.model.vocab_size = embedding_size From 3168e54038ddcc20939a710e792a6608a4b336ac Mon Sep 17 00:00:00 2001 From: Yash Mehan Date: Fri, 27 Feb 2026 05:06:53 +0000 Subject: [PATCH 13/18] put upperbounds on dependencies in pyproject Signed-off-by: Yash Mehan --- pyproject.toml | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 352650f7fd..09095aa276 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,19 +28,19 @@ classifiers=[ dependencies = [ "numpy>=1.26.4,<2.2.0", "accelerate>=1.9.0,<2.0.0", -"transformers==5.2.0", -"torch==2.10.0", -"torchvision>=0.25.0", +"transformers>=5.2.0,<=5.3.0", +"torch>=2.10.0,<2.11.0", +"torchvision<=0.25.1", "sentencepiece>=0.1.99,<0.3", -"tokenizers==0.22.2", +"tokenizers<=0.23.0", "tqdm>=4.66.2,<5.0", "trl>=0.19.1,<0.20.0", "peft>=0.18.1,<0.19.0", "datasets>=4.0.0,<5.0.0", "simpleeval>=0.9.13,<2.0", "pillow>=12.1.1", -"kernels==0.12.1", -"huggingface_hub>=1.3.0" +"kernels>=0.12.1,<0.13.0", +"huggingface_hub>=1.3.0,<1.4.0" ] [project.optional-dependencies] From c507131e697344e8e7da8e6c9bb7ed735b3be3a0 Mon Sep 17 00:00:00 2001 From: Yash Mehan Date: Fri, 27 Feb 2026 05:07:40 +0000 Subject: [PATCH 14/18] fix input_ids Signed-off-by: Yash Mehan --- tuning/data/data_handlers.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tuning/data/data_handlers.py b/tuning/data/data_handlers.py index 327fb40acf..f549f18a24 100644 --- a/tuning/data/data_handlers.py +++ b/tuning/data/data_handlers.py @@ -532,7 +532,7 @@ def tokenize_and_apply_chat_template_with_masking( add_generation_prompt=False, tools=tools, documents=documents, - ) + )["input_ids"] ) # clone labels from input ids @@ -557,7 +557,7 @@ def tokenize_and_apply_chat_template_with_masking( add_generation_prompt=False, tools=tools, documents=documents, - ).shape[1] + )["input_ids"].shape[1] ) # next, we calculate the end index of this non-assistant message if ( @@ -578,7 +578,7 @@ def tokenize_and_apply_chat_template_with_masking( add_generation_prompt=True, tools=tools, documents=documents, - ).shape[1] + )["input_ids"].shape[1] ) else: # for the last message or the message that doesn't follow with @@ -594,7 +594,7 @@ def tokenize_and_apply_chat_template_with_masking( add_generation_prompt=False, tools=tools, documents=documents, - ).shape[1] + )["input_ids"].shape[1] ) # set the label to -100 for the non-assistant part labels[:, message_start_idx:message_end_idx] = -100 From 4f482f1fb467ea9892ee8e37ba62d3145566333d Mon Sep 17 00:00:00 2001 From: Yash Mehan Date: Mon, 2 Mar 2026 07:40:26 +0000 Subject: [PATCH 15/18] linting and formatting Signed-off-by: Yash Mehan --- build/accelerate_launch.py | 2 +- pyproject.toml | 8 ++++---- tests/utils/test_embedding_resize.py | 15 ++++++++------- tuning/data/tokenizer_utils.py | 23 ++++++++++++----------- tuning/sft_trainer.py | 4 +++- 5 files changed, 28 insertions(+), 24 deletions(-) diff --git a/build/accelerate_launch.py b/build/accelerate_launch.py index 5926c4a852..43c6e19e68 100644 --- a/build/accelerate_launch.py +++ b/build/accelerate_launch.py @@ -120,7 +120,7 @@ def main(): logging.error(traceback.format_exc()) # v5: torch.distributed raises ChildFailedError with per-rank exit codes # Check if the root cause was a user error - if hasattr(e, 'failures'): + if hasattr(e, "failures"): root_codes = [f.exitcode for f in e.failures.values()] if any(c == USER_ERROR_EXIT_CODE for c in root_codes): sys.exit(USER_ERROR_EXIT_CODE) diff --git a/pyproject.toml b/pyproject.toml index 09095aa276..dab1a100a3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,9 +28,9 @@ classifiers=[ dependencies = [ "numpy>=1.26.4,<2.2.0", "accelerate>=1.9.0,<2.0.0", -"transformers>=5.2.0,<=5.3.0", -"torch>=2.10.0,<2.11.0", -"torchvision<=0.25.1", +"transformers>=5.2.0,<5.3.0", +"torch>2.7.0,<=2.9.0", +"torchvision<=0.24.0", "sentencepiece>=0.1.99,<0.3", "tokenizers<=0.23.0", "tqdm>=4.66.2,<5.0", @@ -40,7 +40,7 @@ dependencies = [ "simpleeval>=0.9.13,<2.0", "pillow>=12.1.1", "kernels>=0.12.1,<0.13.0", -"huggingface_hub>=1.3.0,<1.4.0" +"huggingface_hub>=1.3.0,<1.4.0", ] [project.optional-dependencies] diff --git a/tests/utils/test_embedding_resize.py b/tests/utils/test_embedding_resize.py index e723b486ba..2ef4513b8f 100644 --- a/tests/utils/test_embedding_resize.py +++ b/tests/utils/test_embedding_resize.py @@ -20,11 +20,9 @@ # Third Party from transformers import ( - AutoModelForCausalLM, - AutoModelForImageTextToText, #AutoModelForVision2Seq was renamed to this in transformers v5 - AutoProcessor, - AutoTokenizer, + AutoModelForImageTextToText, # AutoModelForVision2Seq was renamed to this in transformers v5 ) +from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer import torch # First Party @@ -136,8 +134,9 @@ def test_special_tokens_before_and_after(): special_tokens_dict = {"sep_token": "", "pad_token": ""} addn_spl_tokens_added = ["", "", ""] # for transformers v5: merge existing extra_special_tokens with new ones to prevent replacement - special_tokens_dict["additional_special_tokens"] = list(tokenizer.extra_special_tokens) + addn_spl_tokens_added - + special_tokens_dict["additional_special_tokens"] = ( + list(tokenizer.extra_special_tokens) + addn_spl_tokens_added + ) resize_result = tokenizer_and_embedding_resize( special_tokens_dict=special_tokens_dict, @@ -210,7 +209,9 @@ def test_resize_with_multiple_of(): def test_resize_llama_vision_model(): - model = AutoModelForImageTextToText.from_pretrained(TINY_LLAMA_VISION_MODEL_NAME) # AutoModelForVision2Seq was renamed to AutoModelForImageTextToText in transformers v5 + model = AutoModelForImageTextToText.from_pretrained( + TINY_LLAMA_VISION_MODEL_NAME + ) # AutoModelForVision2Seq was renamed to AutoModelForImageTextToText in transformers v5 processor = AutoProcessor.from_pretrained(TINY_LLAMA_VISION_MODEL_NAME) tokenizer = processor.tokenizer diff --git a/tuning/data/tokenizer_utils.py b/tuning/data/tokenizer_utils.py index eb5ce5fb8d..13faa47d8f 100644 --- a/tuning/data/tokenizer_utils.py +++ b/tuning/data/tokenizer_utils.py @@ -45,21 +45,24 @@ def get_special_tokens_dict( special_tokens_dict = {} if not tokenizer_name_or_path: llama_classes = tuple( - cls for cls in [ + cls + for cls in [ getattr(transformers, "LlamaTokenizer", None), getattr(transformers, "LlamaTokenizerFast", None), - ] if cls is not None + ] + if cls is not None ) is_llama_tokenizer = ( - (bool(llama_classes) and isinstance(tokenizer, llama_classes)) - or "llama" in (getattr(tokenizer, "name_or_path", "") or "").lower() - ) + bool(llama_classes) and isinstance(tokenizer, llama_classes) + ) or "llama" in (getattr(tokenizer, "name_or_path", "") or "").lower() gpt_neox_classes = tuple( - cls for cls in [ + cls + for cls in [ getattr(transformers, "GPTNeoXTokenizerFast", None), getattr(transformers, "GPTNeoXTokenizer", None), - ] if cls is not None + ] + if cls is not None ) if is_llama_tokenizer: @@ -67,9 +70,7 @@ def get_special_tokens_dict( special_tokens_dict["eos_token"] = "" special_tokens_dict["unk_token"] = "" special_tokens_dict["pad_token"] = "" - elif isinstance( - tokenizer, (transformers.GPT2Tokenizer, *gpt_neox_classes) - ): + elif isinstance(tokenizer, (transformers.GPT2Tokenizer, *gpt_neox_classes)): special_tokens_dict["pad_token"] = "" # Add special tokens only when a custom tokenizer is not passed @@ -117,7 +118,7 @@ def tokenizer_and_embedding_resize( dict: Metadata on number of added tokens. """ num_new_tokens = tokenizer.add_special_tokens( - special_tokens_dict=special_tokens_dict, + special_tokens_dict=special_tokens_dict, # replace_additional_special_tokens=False ) embedding_size = int(multiple_of * math.ceil(len(tokenizer) / multiple_of)) diff --git a/tuning/sft_trainer.py b/tuning/sft_trainer.py index 51ab291603..51d23fe165 100644 --- a/tuning/sft_trainer.py +++ b/tuning/sft_trainer.py @@ -28,9 +28,11 @@ from peft import LoraConfig from peft.utils.other import fsdp_auto_wrap_policy from torch.cuda import OutOfMemoryError +from transformers import ( + AutoModelForImageTextToText, # AutoModelForVision2Seq was renamed in transformers v5 +) from transformers import ( AutoModelForCausalLM, - AutoModelForImageTextToText, # AutoModelForVision2Seq was renamed in transformers v5 AutoProcessor, AutoTokenizer, TrainerCallback, From e14d2ce8149af32bb0a0a264f4e50ab83733f595 Mon Sep 17 00:00:00 2001 From: Yash Mehan Date: Fri, 6 Mar 2026 06:51:29 +0000 Subject: [PATCH 16/18] remove multiple excepts Signed-off-by: Yash Mehan --- build/accelerate_launch.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/build/accelerate_launch.py b/build/accelerate_launch.py index 43c6e19e68..968753bfd9 100644 --- a/build/accelerate_launch.py +++ b/build/accelerate_launch.py @@ -110,7 +110,6 @@ def main(): # message to termination log. logging.error(traceback.format_exc()) # The exit code that sft_trainer.py threw is captured in e.returncode - return_code = e.returncode if return_code not in [INTERNAL_ERROR_EXIT_CODE, USER_ERROR_EXIT_CODE]: return_code = INTERNAL_ERROR_EXIT_CODE @@ -126,10 +125,6 @@ def main(): sys.exit(USER_ERROR_EXIT_CODE) write_termination_log(f"Unhandled exception during training. {e}") sys.exit(INTERNAL_ERROR_EXIT_CODE) - except Exception as e: # pylint: disable=broad-except - logging.error(traceback.format_exc()) - write_termination_log(f"Unhandled exception during training. {e}") - sys.exit(INTERNAL_ERROR_EXIT_CODE) peft_method = job_config.get("peft_method") From fd3b97700888fc2585b392cb428b4680c8049c7f Mon Sep 17 00:00:00 2001 From: Yash Mehan Date: Fri, 6 Mar 2026 06:52:08 +0000 Subject: [PATCH 17/18] add messages control to remove in pylintrc Signed-off-by: Yash Mehan --- .pylintrc | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.pylintrc b/.pylintrc index 570205ae40..612fa0e8f2 100644 --- a/.pylintrc +++ b/.pylintrc @@ -447,7 +447,10 @@ disable=raw-checker-failed, duplicate-code, unbalanced-tuple-unpacking, unspecified-encoding, - too-many-lines + too-many-lines, + no-name-in-module, + unexpected-keyword-arg, + unused-argument # Enable the message, report, category or checker with the given id(s). You can # either give multiple identifier separated by comma (,) or put this option From c854702c8422b4cda2f42256429444db0c70c7ae Mon Sep 17 00:00:00 2001 From: Yash Mehan Date: Fri, 6 Mar 2026 16:21:47 +0000 Subject: [PATCH 18/18] upgrade support to trl 0.27+ Signed-off-by: Yash Mehan --- pyproject.toml | 2 +- tests/data/test_data_preprocessing.py | 3 +- tuning/data/collators.py | 245 ++++++++++++++++++++++++ tuning/data/data_preprocessing_utils.py | 3 +- 4 files changed, 248 insertions(+), 5 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index dab1a100a3..27fdaaa7b7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,7 +34,7 @@ dependencies = [ "sentencepiece>=0.1.99,<0.3", "tokenizers<=0.23.0", "tqdm>=4.66.2,<5.0", -"trl>=0.19.1,<0.20.0", +"trl>=0.27.0,<0.29.0", "peft>=0.18.1,<0.19.0", "datasets>=4.0.0,<5.0.0", "simpleeval>=0.9.13,<2.0", diff --git a/tests/data/test_data_preprocessing.py b/tests/data/test_data_preprocessing.py index a1072d2ece..0e7153932d 100644 --- a/tests/data/test_data_preprocessing.py +++ b/tests/data/test_data_preprocessing.py @@ -22,7 +22,6 @@ from datasets import Dataset, DatasetDict, IterableDataset from PIL import Image from transformers import AutoProcessor, AutoTokenizer, DataCollatorForSeq2Seq -from trl import DataCollatorForCompletionOnlyLM import datasets import numpy as np import pyarrow @@ -69,7 +68,7 @@ # Local from tuning.config import configs from tuning.config.acceleration_configs import AttentionAndDistributedPackingConfig -from tuning.data.collators import VisionDataCollator +from tuning.data.collators import DataCollatorForCompletionOnlyLM, VisionDataCollator from tuning.data.data_config import ( DataHandlerConfig, DataPreProcessorConfig, diff --git a/tuning/data/collators.py b/tuning/data/collators.py index c7f63a99cc..b9947e3aba 100644 --- a/tuning/data/collators.py +++ b/tuning/data/collators.py @@ -12,6 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. +# Standard +from typing import Any, Optional, Union +import logging + +# Third Party +from transformers import DataCollatorForLanguageModeling +import numpy as np +import torch + # Local from tuning.data.utils import try_convert_bytes_dict_to_pil @@ -91,3 +100,239 @@ def __call__(self, features): batch["labels"] = labels return batch + + +class DataCollatorForCompletionOnlyLM(DataCollatorForLanguageModeling): + """ + Data collator used for completion tasks. + It ensures that all the tokens of the labels + are set to an 'ignore_index' + when they do not come from the assistant. + This ensure that the loss is only calculated on the completion made by + the assistant. + + Args: + response_template (`Union[str, list[int]]`): + the template form that indicates the + start of the response, typically + something like '### Response:\n'. It + can also be passed as tokenized ids, + which can be useful when using a tokenizer + that encodes the response + differently if it does not have proper context. + instruction_template (`Union[str, list[int]]`): + the template form that indicates the start + of the human instruction, typically + something like '### + Human:\n'. Useful for assistant-style + conversation datasets. It can also be passed + as tokenized ids. + mlm (`bool`, *optional*, defaults to `False`): Whether + to use masked language modeling in the underlying + `DataCollatorForLanguageModeling` class. + Note that this option currently has no effect but is present + for flexibility and backwards-compatibility. + ignore_index (`int`, *optional*, defaults to `-100`): + The index to use to ignore the initial tokens with + """ + + def __init__( + self, + *args, + response_template: Union[str, list[int]], + instruction_template: Optional[Union[str, list[int]]] = None, + mlm: bool = False, + ignore_index: int = -100, + padding_free: bool = False, + **kwargs, + ): + super().__init__(*args, mlm=mlm, **kwargs) + + self.instruction_template = instruction_template + if isinstance(instruction_template, str): + # The user provides a string, must tokenize + self.instruction_token_ids = self.tokenizer.encode( + self.instruction_template, add_special_tokens=False + ) + else: + # The user already provides the token ids + self.instruction_token_ids = instruction_template + + self.response_template = response_template + if isinstance(response_template, str): + # The user provides a string, must tokenize + self.response_token_ids = self.tokenizer.encode( + self.response_template, add_special_tokens=False + ) + else: + # The user already provides the token ids + self.response_token_ids = response_template + + if ( + not self.mlm + and self.instruction_template + and self.tokenizer.pad_token_id == self.tokenizer.eos_token_id + ): + logging.warning( + "The pad_token_id and eos_token_id values " + "of this tokenizer are identical. " + "If you are planning for multi-turn training, " + "it can result in the model continuously generating " + "questions and answers without eos token. " + "To avoid this, set the pad_token_id to a different value.", + ) + + self.ignore_index = ignore_index + self.padding_free = padding_free + + def torch_call( + self, examples: list[Union[list[int], Any, dict[str, Any]]] + ) -> dict[str, Any]: + batch = super().torch_call(examples) + + if self.instruction_template is None: + for i in range(len(examples)): + response_token_ids_start_idx = None + + for idx in np.where(batch["labels"][i] == self.response_token_ids[0])[ + 0 + ]: + # `response_token_ids` is + # `'### Response:\n'`, here we are just making sure + # that the token IDs match + if ( + self.response_token_ids + == batch["labels"][i][ + idx : idx + len(self.response_token_ids) + ].tolist() + ): + response_token_ids_start_idx = idx + + if response_token_ids_start_idx is None: + logging.warning( + "Could not find response key %s in the following instance: " + "%s. This instance will be ignored in loss " + "calculation. Note, if this happens often, " + "consider increasing the `max_length`.", + self.response_template, + self.tokenizer.decode(batch["input_ids"][i]), + ) + batch["labels"][i, :] = self.ignore_index + else: + response_token_ids_end_idx = response_token_ids_start_idx + len( + self.response_token_ids + ) + + # Make pytorch loss function ignore all + # tokens up through the end of the response key + batch["labels"][i, :response_token_ids_end_idx] = self.ignore_index + + else: + for i in range(len(examples)): + response_token_ids_idxs = [] + human_token_ids_idxs = [] + + for assistant_idx in np.where( + batch["labels"][i] == self.response_token_ids[0] + )[0]: + # find the indexes of the start of a response. + if ( + self.response_token_ids + == batch["labels"][i][ + assistant_idx : assistant_idx + len(self.response_token_ids) + ].tolist() + ): + response_token_ids_idxs.append( + assistant_idx + len(self.response_token_ids) + ) + + if len(response_token_ids_idxs) == 0: + logging.warning( + "Could not find response key %s in the following instance: " + "%s. This instance will be ignored in loss " + "calculation. Note, if this happens often, " + "consider increasing the `max_length`.", + self.response_template, + self.tokenizer.decode(batch["input_ids"][i]), + ) + batch["labels"][i, :] = self.ignore_index + + human_token_ids = self.instruction_token_ids + for human_idx in np.where(batch["labels"][i] == human_token_ids[0])[0]: + # find the indexes of the start of a human answer. + if ( + human_token_ids + == batch["labels"][i][ + human_idx : human_idx + len(human_token_ids) + ].tolist() + ): + human_token_ids_idxs.append(human_idx) + + if len(human_token_ids_idxs) == 0: + logging.warning( + "Could not find instruction key `%s` in the following instance: " + "%s. This instance will be ignored in loss " + "calculation. Note, if this happens often, " + "consider increasing the `max_length`.", + self.instruction_template, + self.tokenizer.decode(batch["input_ids"][i]), + ) + batch["labels"][i, :] = self.ignore_index + + if ( + len(human_token_ids_idxs) > 0 + and len(response_token_ids_idxs) > 0 + and human_token_ids_idxs[0] > response_token_ids_idxs[0] + ): + human_token_ids_idxs = [0] + human_token_ids_idxs + + for idx, (start, end) in enumerate( + zip(human_token_ids_idxs, response_token_ids_idxs) + ): + # Make pytorch loss function ignore all non response tokens + if idx != 0: + batch["labels"][i, start:end] = self.ignore_index + else: + batch["labels"][i, :end] = self.ignore_index + + if len(response_token_ids_idxs) < len(human_token_ids_idxs): + batch["labels"][i, human_token_ids_idxs[-1] :] = self.ignore_index + + if self.padding_free: + # remove padding, `attention_mask` and add `position_ids` + attn_mask = batch.pop("attention_mask") + batch["input_ids"] = batch["input_ids"][attn_mask.bool()].unsqueeze(0) + batch["position_ids"] = ( + attn_mask.cumsum(1)[attn_mask.bool()].unsqueeze(0) - 1 + ) + batch["labels"] = batch["labels"][attn_mask.bool()].unsqueeze(0) + batch["labels"][batch["position_ids"] == 0] = self.ignore_index + + # Calculate cumulative sequence lengths for queries and + # keys to prevent graph breaks during further computations. + flattened_position_ids = batch["position_ids"].flatten() + indices_q = torch.arange( + flattened_position_ids.size(0), + device=flattened_position_ids.device, + dtype=torch.int32, + ) + batch["cu_seq_lens_q"] = torch.cat( + ( + indices_q[flattened_position_ids == 0], + torch.tensor( + flattened_position_ids.size(), + device=flattened_position_ids.device, + dtype=torch.int32, + ), + ) + ).unsqueeze(0) + batch["cu_seq_lens_k"] = batch["cu_seq_lens_q"] + + # Determine maximum sequence lengths to + # prevent graph breaks during further computations. + batch["max_length_k"] = torch.tensor( + [flattened_position_ids.max().item() + 1] + ) + batch["max_length_q"] = batch["max_length_k"] + + return batch diff --git a/tuning/data/data_preprocessing_utils.py b/tuning/data/data_preprocessing_utils.py index 04d0d7a2d4..b43081864d 100644 --- a/tuning/data/data_preprocessing_utils.py +++ b/tuning/data/data_preprocessing_utils.py @@ -22,11 +22,10 @@ DataCollatorForSeq2Seq, LlavaProcessor, ) -from trl import DataCollatorForCompletionOnlyLM # Local from tuning.config import configs -from tuning.data.collators import VisionDataCollator +from tuning.data.collators import DataCollatorForCompletionOnlyLM, VisionDataCollator logger = logging.getLogger(__name__)