Causal Language Model Trainer

This is the node that actually does the work in the LLM School pack (heshengtao/comfyui_LLM_schools), the fine-tuning companion to the same author's far more famous comfyui_LLM_party. Instead of calling an API for a chat completion, it runs a full training loop on a causal language model from the transformers library - GPT-2 by default, but any AutoModelForCausalLM on the Hub works, including ones you've downloaded locally. If you've ever wanted to fine-tune an LLM but bounced off a wall of notebooks and shell commands, this is the on-ramp that lives inside your graph.

How it actually works (read this before you run it)

The big surprise: the node doesn't train in ComfyUI's process. It builds a command line and fires off python train_core.py as a subprocess - a brand-new console window on Windows, Terminal.app on macOS, and on Linux a background process whose PID it prints. That means the log output is just a "training started" message, not real progress, and on Linux you'll be watching top or ps to see anything happen. It also means the spawned python needs to be an environment with torch, transformers, peft, and datasets installed - if ComfyUI runs in a venv and python on your PATH is a different one, the new process crashes with ModuleNotFoundError before training starts. Keep your python environments straight or this bites you immediately.

Second surprise, in the training code itself: the preprocessing is SQuAD-style extractive QA, not generic text. train_core.py expects columns named question, context, and answers (with answer_start/text), computes answer span positions, and collates with DataCollatorForLanguageModeling(mlm=False). The pack's default dataset, rajpurkar/squad_v2, matches this exactly. Feed it a plain text dataset and it'll KeyError. So despite the name, think of this as "fine-tune a causal LM on question-answering data" rather than "train on your corpus."

The inputs that matter

• split_datapaths - path to the train/val/test folder from the pack's split_dataset node. Required, no default.
• model_name_or_path - any HF causal LM repo id or local dir. gpt2 default is honest about size: a full fine-tune of anything bigger will eat your VRAM.
• dtype - float32 / float16 / bfloat16 / int8 / int4. The int options use BitsAndBytesConfig quantization, which is how you squeeze a bigger model onto the card.
• fine_tuning_method - full_fine_tuning, lora, adapter, p-tuning, prompt_tuning, prefix-tuning, IA3.
• args (from LLM_Arguments) and optional peft_args (from the per-method argument nodes).

Output is a single log string - wire it to a Show Text node so you at least get the "started" confirmation.

Where people get burned

• The default method crashes. In the shipped train_core.py, full_fine_tuning is not handled in the PEFT dispatch - it falls through to raise ValueError("Unsupported fine-tuning method: full_fine_tuning"). The README claims full fine-tuning works, but this build of the code doesn't. Pick lora (or another listed method) or patch the dispatch.
• Training ends with a KeyError. After training, trainer.evaluate() returns a plain dict, but the code reads results['accuracy'] and results['other_metrics'], which don't exist → crash after the run finishes. Checkpoints are already saved by then, so it's cosmetic - but it looks alarming.
• It re-runs every queue. IS_CHANGED returns a fresh timestamp hash, so the node is never cached. Fine for a one-off; annoying if it sits in a big workflow.

Install

ComfyUI Manager, search "comfyui_LLM_schools" - or:

cd ComfyUI/custom_nodes
git clone https://github.com/heshengtao/comfyui_LLM_schools

then restart and pip install -r requirements.txt (that's huggingface_hub, datasets, transformers, peft). It's a niche, early pack with almost no community footprint, so expect to be your own tech support - but if you want a first LLM fine-tune from inside a graph, nothing else in ComfyUI does this end to end.