Add decoder-only LLM fine-tuning path (LoRA/QLoRA)#1109
Conversation
- Add LLMConfig/LLMModel for decoder-only causal LMs (full/LoRA/QLoRA), dispatched from config_utils on `model_type: llm` - Reuse the existing preprocess/test/translate/experiment pipeline via the Config/NMTModel interfaces; chat-template prompts with prompt-only label masking - Lift shared checkpoint resolution, effective-config, and inference-cache logic into the Config/NMTModel base classes - Rename hugging_face_config.py -> seq2seq_config.py (Seq2SeqConfig/Seq2SeqNMTModel) - Add bitsandbytes as an optional `llm` extra for QLoRA - Add smoke tests, fixtures, and unit tests for the LLM path Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
This PR adds general support for LLMs. There is a little bit more work that is needed to support TranslateGemma and some other models. A newer version of transformers is needed. TranslateGemma is a vision language model, i.e. it supports both text and image input. These means we will need to add support for |
benjaminking
left a comment
There was a problem hiding this comment.
I've made a number of comments with the assumption that it should be fairly quick to unleash Claude on them. There are several that wouldn't need to be addressed immediately if they turn out to be more time consuming.
@benjaminking partially reviewed 16 files and all commit messages, and made 11 comments.
Reviewable status: all files reviewed, 10 unresolved discussions (waiting on ddaspit).
silnlp/nmt/llm_config.py line 189 at r1 (raw file):
num_devices: int = 1, clearml_queue: Optional[str] = None, pretrained_model_provider_factory: "CausalLMProviderFactory" = None, # type: ignore[assignment]
Could you use a type hint of Optional["CausalLMProviderFactory"] here?
silnlp/nmt/llm_config.py line 213 at r1 (raw file):
@property def train_src_iso(self) -> str: return self.default_test_src_iso or (next(iter(self.src_isos)) if len(self.src_isos) > 0 else "")
Does this support multiple source languages or multiple target languages? If not, we should probably refuse to process any config set up that way with an LLM.
silnlp/nmt/llm_config.py line 221 at r1 (raw file):
def build_prompt_messages( self, source: str, src_lang: str, trg_lang: str, target: Optional[str] = None ) -> List[Dict[str, str]]:
I'd lean toward representing the prompt messages with a class, so that you can encapsulate all the object creation logic from this method in the class (along with _fold_system_message) and have more semantically meaningful type hints that use the class name. The class would unfortunately still need to dump out a Dict[str, str] to be compatible with the apply_chat_template method.
silnlp/nmt/llm_config.py line 502 at r1 (raw file):
trainer.save_state() def _apply_finetuning(self, model: PreTrainedModel) -> PreTrainedModel:
Could rename this method to make it clearer that the fine tuning is not happening here, but it's just applying the fine tuning configuration?
silnlp/nmt/llm_config.py line 537 at r1 (raw file):
open(trg_path, "r", encoding="utf-8-sig") as trg_file, ): for src_line, trg_line in zip(src_file, trg_file):
I believe we use Machine's Corpus class for similar functionality elsewhere. Does it make sense to use it here, to minimize the number of places we need to update if the file format changes?
silnlp/nmt/llm_config.py line 544 at r1 (raw file):
return Dataset.from_dict({"src": sources, "trg": targets}) def _create_training_arguments(self) -> TrainingArguments:
This is almost the same as the version in Seq2SeqNMTModel. Is it possible to pull this up into the parent class?
silnlp/nmt/llm_config.py line 673 at r1 (raw file):
"pad_token_id": tokenizer.pad_token_id, } if num_return_sequences > 1 or infer.get("do_sample"):
This logic doesn't track for me. The number of return sequences should be independent of whether we want to sample. We can return multiple sequences for both beam search and sampling. And if we are sampling, then num_beams should not be set.
silnlp/nmt/llm_config.py line 701 at r1 (raw file):
transition_scores = None if save_confidences and getattr(output, "scores", None) is not None: transition_scores = model.compute_transition_scores(
If num_beams > 1, we need to include the beam indices in the call to compute_transition_scores.
silnlp/nmt/llm_config.py line 718 at r1 (raw file):
if len(valid) > 0: sequence_score = sum(valid) / len(valid) # tokens=["", text] so join_tokens_for_test_file() (which drops tokens[0])
I think it would be better to update join_tokens_for_test_file and condition it on the model type so that it doesn't automatically drop the first token when it's working with LLM output.
tests/smoke_tests/test_llm_config_units.py at r1 (raw file):
Could you create a separate folder for unit tests and move these there? These seem more like unit tests than smoke tests.
…rect auto model class when loading them
Support image+text LLMs
model_type: llmllmextra for QLoRAThis change is