Fine-Tune a Quantized Llama on GSM8K with PEFT LoftQ

Quantized models are cheap to serve, but fine-tuning them has traditionally been a hostile process: the quantization error baked into the weights becomes part of what you are optimizing, and standard LoRA initialization does nothing to repair it. LoftQ changes that calculation. Instead of initializing the low-rank adapter randomly, LoftQ solves for an adapter that approximates the difference between the original weights and their quantized versions, then begins training from that repaired state. For GSM8K, exact numeric answers matter: the example script reports exact-match accuracy, so a model that drifts in intermediate arithmetic can produce a plausible-looking but wrong final number.

This guide walks through the Hugging Face PEFT example that fine-tunes a quantized Llama model on GSM8K with LoftQ. It is based on train_gsm8k_llama.py in the huggingface/peft repository, licensed under Apache-2.0. Code excerpts below are reproduced verbatim from that source; the surrounding explanation is original.

Prerequisites

The source establishes a handful of hard requirements. You need a Python environment with torch, transformers, accelerate, datasets, bitsandbytes, and peft installed. The script checks for datasets>=1.8.0. A CUDA-capable GPU is expected, but the source does not declare a minimum VRAM. The memory levers it exposes are 4-bit NF4 loading, low_cpu_mem_usage=True, per-device batch size, and gradient accumulation. The defaults that affect the run are a per-device batch size of 8, three epochs, maximum source length 128, maximum target length 128, and learning rate 5e-5. The source does not report wall-clock runtimes, VRAM usage, or cloud cost; those depend on the base model, GPU, and chosen split. A Hugging Face account and token are required only if you pass --push_to_hub. The script also sends argument, Python, and PyTorch version telemetry via send_example_telemetry.

Step 1: Load and split GSM8K

GSM8K is a math word problem dataset with question and answer columns. The script loads it from the Hub and, if no validation split is present, carves one out of the training set using the percentage supplied by --validation_split_percentage. The default is 5, so the first 5% of train becomes validation and the remaining 95% stays as train. The same else branch supports CSV, JSON, or TXT files; a .txt file is treated as the text loader and can preserve line breaks unless --no_keep_linebreaks is passed.

    if args.dataset_name is not None:
        # Downloading and loading a dataset from the hub.
        raw_datasets = load_dataset(args.dataset_name, args.dataset_config_name)
        if "validation" not in raw_datasets.keys():
            raw_datasets["validation"] = load_dataset(
                args.dataset_name,
                args.dataset_config_name,
                split=f"train[:{args.validation_split_percentage}%]",
            )
            raw_datasets["train"] = load_dataset(
                args.dataset_name,
                args.dataset_config_name,
                split=f"train[{args.validation_split_percentage}%:]",
            )
    else:
        data_files = {}
        dataset_args = {}
        if args.train_file is not None:
            data_files["train"] = args.train_file
        if args.validation_file is not None:
            data_files["validation"] = args.validation_file
        extension = args.train_file.split(".")[-1]
        if extension == "txt":
            extension = "text"
            dataset_args["keep_linebreaks"] = not args.no_keep_linebreaks
        raw_datasets = load_dataset(extension, data_files=data_files, **dataset_args)
        # If no validation data is there, validation_split_percentage will be used to divide the dataset.
        if "validation" not in raw_datasets.keys():
            raw_datasets["validation"] = load_dataset(
                extension,
                data_files=data_files,
                split=f"train[:{args.validation_split_percentage}%]",
                **dataset_args,
            )
            raw_datasets["train"] = load_dataset(
                extension,
                data_files=data_files,
                split=f"train[{args.validation_split_percentage}%:]",
                **dataset_args,
            )

Step 2: Load the 4-bit base model and attach the LoftQ adapter

The tokenizer is configured for causal language modeling with left padding and left truncation. Left padding keeps the prompt at the right side of the sequence, which matters for generation and batched evaluation. The model is loaded in 4-bit NF4 with low_cpu_mem_usage=True, so parameters are materialized only as pretrained weights are loaded rather than allocating the full model in memory first. Double quantization is disabled, and the compute dtype is taken from the model config.

    ##########################
    #        Tokenizer       #
    ##########################
    tokenizer.pad_token_id = 0  # unk. we want this to be different from the eos token
    tokenizer.padding_side = "left"  # Allow batched inference
    tokenizer.truncation_side = "left"

    if args.model_name_or_path:
        model = AutoModelForCausalLM.from_pretrained(
            args.model_name_or_path,
            from_tf=bool(".ckpt" in args.model_name_or_path),
            config=config,
            low_cpu_mem_usage=True,
            quantization_config=BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_use_double_quant=False,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_compute_dtype=config.dtype,
            ),
        )
    else:
        logger.info("Training new model from scratch")
        model = AutoModelForCausalLM.from_config(config, trust_remote_code=args.trust_remote_code)

The next block is the LoftQ-specific part of the workflow. When --adapter_name_or_path is None, the script loads a precomputed LoftQ initialization from a loftq_init subfolder inside the model repository. Passing a path instead loads an already-trained adapter for evaluation or continued training. In both cases is_trainable=True tells PEFT to keep base model weights frozen and only train the adapter.

    if args.adapter_name_or_path is None:
        model = PeftModel.from_pretrained(model, args.model_name_or_path, subfolder="loftq_init", is_trainable=True)
    else:
        model = PeftModel.from_pretrained(model, args.adapter_name_or_path, is_trainable=True)
    model.print_trainable_parameters()

    embedding_size = model.get_input_embeddings().weight.shape[0]
    if len(tokenizer) > embedding_size:
        model.resize_token_embeddings(len(tokenizer))

The table below summarizes the key quantization and adapter decisions that affect memory and behavior.

Step 3: Preprocess GSM8K into a causal LM format

The preprocessing block adapts GSM8K for generation. It replaces the dataset's #### delimiter with the phrase The final answer is, then concatenates the question, a task prompt that asks the model to think step by step and answer the final number, and the full reasoning target. The source and target column names are hardcoded to question and answer. During training, tokens are padded or batched dynamically according to --pad_to_max_length. With max-length padding, labels are a deep copy of the input IDs; the prefix before the target tokens is set to -100, so the loss only covers the rationale and final answer. The test preprocessing supplies the source question and prompt, while labels are the gold answer tokens for scoring.

    column_names = raw_datasets["train"].column_names

    # Get the column names for source/target.
    source_column, target_column = "question", "answer"

    # Temporarily set max_target_length for training.
    padding = "max_length" if args.pad_to_max_length else False
    task_prompt = "\nAnswer the above question. First think step by step and then answer the final number.\n"

    def prompt_process(sent_1, sent_2, prompt_1="", prompt_2="", prompt_3=""):
        sent_2 = sent_2.replace("####", "The final answer is")
        return prompt_1 + sent_1 + prompt_2 + sent_2 + prompt_3

    def preprocess_function_train(examples):
        sources = examples[source_column]
        targets = examples[target_column]

        inputs = [prompt_process(source, target, prompt_2=task_prompt) for (source, target) in zip(sources, targets)]

        model_inputs = tokenizer(
            inputs,
            max_length=args.max_source_length + args.max_target_length,
            padding=padding,
            truncation=True,
            return_tensors="pt",
        )

        labels = copy.deepcopy(model_inputs)

        # If we are padding here, replace all tokenizer.pad_token_id in the labels by -100 when we want to ignore
        # padding in the loss.
        if padding == "max_length" and args.ignore_pad_token_for_loss:
            # get the length of the target tokens. -1 to kick out the <BOS> token
            target_tokens = tokenizer(targets, padding=False)
            target_len = [len(label) - 1 for label in target_tokens["input_ids"]]

            # don't calculate the loss from source and padding (left padding)
            for i in range(len(labels["input_ids"])):
                labels["input_ids"][i, : -target_len[i]] = -100

        model_inputs["labels"] = labels["input_ids"]
        return model_inputs

    def preprocess_function_test(examples):
        sources = examples[source_column]
        labels = examples[target_column]

        inputs = [source + task_prompt for source in sources]

        model_inputs = tokenizer(inputs, max_length=args.max_source_length, padding=padding, truncation=True)
        labels = tokenizer(labels, max_length=args.max_target_length, padding=padding, truncation=True)

        model_inputs["labels"] = labels["input_ids"]

        return model_inputs

Step 4: Configure the optimizer and run the training loop

Because the base model is frozen, the optimizer should only receive LoRA parameters that require gradients. The first parameter group matches names containing "lora" while excluding bias and layer norm weights, and applies the configured weight decay. The second group catches bias and layer norm parameters. The scheduler is created after the number of update steps is calculated from the dataloader length and gradient accumulation steps.

    no_decay = ["bias", "layer_norm.weight"]
    optimizer_grouped_parameters = [
        {
            "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay) and "lora" in n],
            "weight_decay": args.weight_decay,
        },
        {
            "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
            "weight_decay": 0.0,
        },
    ]
    optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=args.learning_rate)

    overrode_max_train_steps = False
    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
    if args.max_train_steps is None:
        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
        overrode_max_train_steps = True

    lr_scheduler = get_scheduler(
        name=args.lr_scheduler_type,
        optimizer=optimizer,
        num_warmup_steps=args.num_warmup_steps * args.gradient_accumulation_steps,
        num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
    )

The training loop itself uses Accelerator's accumulate context so gradient accumulation happens transparently. The model is placed in training mode, loss is computed from the already-attached labels, and the optimizer updates only trainable LoRA parameters. The script prints the loss whenever the completed step count is not a multiple of 50, which produces a periodic but slightly irregular log.

    for epoch in range(starting_epoch, args.num_train_epochs):
        model.train()
        if args.with_tracking:
            total_loss = 0
        if args.resume_from_checkpoint and epoch == starting_epoch and resume_step is not None:
            active_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step)
        else:
            active_dataloader = train_dataloader
        for step, batch in enumerate(active_dataloader):
            with accelerator.accumulate(model):
                outputs = model(**batch)
                loss = outputs.loss
                if args.with_tracking:
                    total_loss += loss.detach().float()
                accelerator.backward(loss)
                if completed_steps % 50:
                    accelerator.print(f"Epoch: {epoch} | Step: {completed_steps} | Loss: {loss}")
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

            if accelerator.sync_gradients:
                progress_bar.update(1)
                completed_steps += 1

            if isinstance(checkpointing_steps, int):
                if completed_steps % checkpointing_steps == 0:
                    output_dir = f"step_{completed_steps}"
                    if args.output_dir is not None:
                        output_dir = os.path.join(args.output_dir, output_dir)
                    accelerator.save_state(output_dir)
            if completed_steps >= args.max_train_steps:
                break

Step 5: Generate answers and compute exact-match accuracy

At the end of each epoch, the script switches to evaluation. Generation uses sampling with max_new_tokens set to the target length, and the default temperature is 0.8, top-k is 40, and top-p is 0.95. The script extracts the newly generated tokens by slicing off the first args.max_source_length positions, then decodes both predictions and labels. Accuracy is not a probabilistic metric: it is the fraction of examples where the extracted numeric answer exactly equals the gold numeric answer.

        model.eval()
        gen_kwargs = {
            "max_new_tokens": args.max_target_length,
            "temperature": args.temperature,
            "top_k": args.k,
            "top_p": args.p,
            "do_sample": True,
        }
        ans_pred_list = []
        ans_gold_list = []
        for step, batch in enumerate(eval_dataloader):
            with torch.no_grad():
                gen_kwargs["input_ids"] = batch["input_ids"]
                gen_kwargs["attention_mask"] = batch["attention_mask"]
                generated_tokens = accelerator.unwrap_model(model).generate(**gen_kwargs)

            pred_tokens = generated_tokens[:, args.max_source_length :]
            pred_tokens = accelerator.pad_across_processes(pred_tokens, dim=1, pad_index=tokenizer.pad_token_id)
            gold_tokens = batch["labels"]

            if not args.pad_to_max_length:
                gold_tokens = accelerator.pad_across_processes(
                    batch["labels"], dim=1, pad_index=tokenizer.pad_token_id
                )

            pred_tokens, gold_tokens = accelerator.gather_for_metrics((pred_tokens, gold_tokens))
            pred_tokens, gold_tokens = pred_tokens.cpu().numpy(), gold_tokens.cpu().numpy()

            if isinstance(pred_tokens, tuple):
                pred_tokens = pred_tokens[0]
            decoded_pred = tokenizer.batch_decode(pred_tokens, skip_special_tokens=True)
            decoded_gold = tokenizer.batch_decode(gold_tokens, skip_special_tokens=True)

            accelerator.print(decoded_pred)
            ans_pred_list += [extract_answer_number(sentence_pred) for sentence_pred in decoded_pred]
            ans_gold_list += [extract_answer_number(sentence_gold) for sentence_gold in decoded_gold]

        accelerator.print(ans_pred_list)
        accelerator.print(ans_gold_list)
        accuracy = compute_accuracy(ans_gold_list, ans_pred_list)

The numeric extraction logic matters more than it may appear. It strips commas, uses a regex to find all numbers, then tries to take the number immediately following the string "The final answer is ". If that phrase is missing, it falls back to the last numeric token in the generation. If no number is found, it returns infinity, which will never match a gold number.

PATTERN_NUMBER = re.compile(r"-?\d+\.?\d*")


def extract_answer_number(sentence: str) -> float:
    sentence = sentence.replace(",", "")
    pred = PATTERN_NUMBER.findall(sentence)
    if not pred:
        return float("inf")
    segment = sentence.split("The final answer is ")
    if len(segment) > 1:
        pred_answer = segment[1]
        pred_answer = PATTERN_NUMBER.findall(pred_answer)
        if len(pred_answer) > 0:
            pred_answer = pred_answer[0]
        else:
            pred_answer = float(pred[-1])
    else:
        pred_answer = float(pred[-1])

    if isinstance(pred_answer, str):
        try:
            pred_answer = float(pred_answer)
        except ValueError:
            pred_answer = float("inf")
    return pred_answer


def compute_accuracy(pred: list, gold: list):
    acc = 0.0
    for p, g in zip(pred, gold):
        if p == g:
            acc += 1

    return acc / len(pred)

What to watch out for

Use --pad_to_max_length for this script. The label-masking logic only runs when padding == "max_length". With dynamic padding, the DataLoader collator still inserts pad tokens, but they are not converted to -100 in the labels. The model then tries to predict pad tokens at the start of left-padded sequences, which distorts the loss. This is one of the most consequential failure modes in the example.

The model repository must already contain loftq_init. When --adapter_name_or_path is None, the script calls PeftModel.from_pretrained(model, args.model_name_or_path, subfolder="loftq_init"). A standard Llama checkpoint without that subfolder will error immediately. Generate and store the LoftQ initialization first, or point --adapter_name_or_path at an existing adapter.

The tokenizer pad token is forced to 0. For Llama, token 0 is also the unknown token, so pad and unk are conflated. This works for the script's left-padding setup, but it is not a semantically clean pad token. If you change tokenizers, verify that the pad index and label masking still align.

bnb_4bit_compute_dtype=config.dtype can silently slow you down. Many model configs default to float32. On Ampere and later GPUs, float32 dequantized compute reduces the speed benefit of 4-bit training. Override the config dtype or set the model config dtype to bfloat16 before loading if your hardware supports it.

Sampling introduces variance into exact-match accuracy. The evaluation uses do_sample=True with temperature 0.8, top-k 40, and top-p 0.95. GSM8K accuracy is sensitive to sampling noise and to the extraction regex. Run evaluation multiple times or lower the temperature if you need a stable report.

Checkpoint resume logic expects a specific folder name. The script distinguishes epoch and step checkpoints by parsing epoch_{i} or step_{i} from the directory name. If you move or rename a checkpoint directory, resume may start from the wrong step or fail to restore the dataloader position.

The optimizer group may include frozen base parameters. The second group does not filter for "lora" in the name, so it can include frozen bias and layer norm parameters from the base model if PEFT exposes them in named_parameters(). They receive no gradients, but they can occupy a small amount of optimizer state. If you are memory-constrained, filter by p.requires_grad too.

Where to go next

Once you have a trained adapter, the natural next step is to evaluate it outside the training script with a larger GSM8K test set or a different arithmetic benchmark. The same adapter can be served alongside a 4-bit base model using vLLM, TGI, or a custom inference loop that loads the PEFT adapter on top of the quantized model. If you are comparing candidate base models for reasoning adaptation, our coverage of Mistral Large 4.1T examines a larger model family that may change the cost/performance calculus. For deeper LoftQ theory and quantization-aware initialization details, see the PEFT documentation and the original LoftQ paper.