Post-training

Post-training starts from a pretrained model and adapts it with your data: supervised fine-tuning (SFT), parameter-efficient methods such as LoRA and QLoRA, or preference tuning such as DPO. The training code is usually a library call, such as TRL’s SFTTrainer or the Hugging Face Trainer. Flyte runs the stages around it: getting the base model and data onto the GPU, keeping the run recoverable, comparing the result with the base model, and handing the weights to serving.

Define the environments

Fine-tuning needs one GPU environment for training and evaluation, and a CPU environment for data prep and the driver:

import os

main_img = flyte.Image.from_uv_script(__file__, name="llm-fine-tuning-lora-qlora", pre=True)

gpu_env = flyte.TaskEnvironment(
    name="llm-fine-tuning-lora-qlora-gpu",
    image=main_img,
    resources=flyte.Resources(cpu=4, memory="24Gi", gpu=1),
    secrets=[flyte.Secret(key="huggingface-token", as_env_var="HF_TOKEN")],
)

cpu_env = flyte.TaskEnvironment(
    name="llm-fine-tuning-lora-qlora-cpu",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="8Gi"),
    depends_on=[gpu_env],
)

HF_TOKEN = os.environ.get("HF_TOKEN")

• flyte.Image.from_uv_script builds the image from the script’s own PEP 723 dependency header, so the image and your local environment install the same dependency list.
• The Hugging Face token comes from a secret and is set as HF_TOKEN, which transformers and datasets read automatically.
• depends_on=[gpu_env] lets a driver in the CPU environment call the GPU tasks.

A single GPU is enough for LoRA on models of a few billion parameters. For full fine-tuning of a larger model, or fine-tuning with DeepSpeed, use a clustered environment.

Structure the pipeline

A fine-tuning pipeline has the same stages whatever the method:

1. Prepare the data in a CPU task with cache="auto", and return it as a flyte.io.Dir. Later runs with a different learning rate or LoRA rank skip this step.
2. Train in a GPU task with report=True, and return the adapter or model as a flyte.io.Dir. The directory goes to your object store, and the next task gets a reference to it.
3. Evaluate the base model and the fine-tuned model on the same held-out set, and report both scores. A fine-tuning run that doesn’t beat the base model is a result too.
4. Serve or publish the model directory.

Because each step is a task, switching methods is a change to one task’s arguments, not to the pipeline. Comparing LoRA with full fine-tuning is two calls to the same pipeline, run side by side with asyncio.gather.

Avoid downloading the base model every run

A large base model takes minutes to download from Hugging Face, on every GPU task that loads it. Copy it into your object store once and pass the stored copy to each task. See Prefetching models.

Resume from the last checkpoint

GPU nodes get preempted, and a fine-tuning run that restarts from step 0 wastes the hours before the failure. Connect the trainer’s checkpoints to flyte.Checkpoint with a callback:

class FlyteTrainerCheckpointCallback(TrainerCallback):
    """Mirror the Trainer's `output_dir` to the Flyte checkpoint after each epoch."""

    def __init__(self, checkpoint: flyte.Checkpoint, output_dir: pathlib.Path) -> None:
        self._checkpoint = checkpoint
        self._output_dir = output_dir

    def on_epoch_end(self, args, state, control, **kwargs) -> None:
        # Trainer callbacks are synchronous, so use save_sync
        self._checkpoint.save_sync(self._output_dir)

Then load the previous attempt’s checkpoint at the start of the task and pass it to trainer.train:

@env.task(retries=RETRIES)
def train_transformers(max_epochs: int = 10) -> float:
    checkpoint = flyte.ctx().checkpoint

    ckpt_dir = pathlib.Path("hf_trainer")
    ckpt_dir.mkdir(parents=True, exist_ok=True)

    # Restore the previous attempt's checkpoint tree and find the last HF checkpoint.
    hf_resume = None
    prev = checkpoint.load_sync()
    if prev:
        hf_resume = get_last_checkpoint(str(prev))

    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID, num_labels=2)

    args = TrainingArguments(
        output_dir=str(ckpt_dir),
        num_train_epochs=max_epochs,
        per_device_train_batch_size=4,
        save_strategy="epoch",
        save_total_limit=2,
        logging_steps=1,
        report_to="none",
        seed=42,
        use_cpu=True,
    )

    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=ToyTextDataset(tokenizer),
        data_collator=DataCollatorWithPadding(tokenizer),
        callbacks=[FlyteTrainerCheckpointCallback(checkpoint, ckpt_dir)],
    )
    trainer.train(resume_from_checkpoint=hf_resume)

    model.eval()
    device = next(model.parameters()).device
    with torch.no_grad():
        batch = tokenizer(
            "classification example for inference",
            return_tensors="pt",
            truncation=True,
            max_length=32,
            padding="max_length",
        )
        batch = {k: v.to(device) for k, v in batch.items()}
        logits = model(**batch).logits
        return float(logits[0, 1].item())

With retries set on the task, a failed attempt restarts from the last saved epoch. The same pattern works for PyTorch Lightning, plain PyTorch loops, and Unsloth. The Brain tumor MRI classification tutorial does it for Lightning. See Intra-task checkpoints.

Watch training while it runs

report=True gives the task a report tab in the UI. Publish loss curves, learning-rate schedules, and sample generations with flyte.report.replace.aio(..., do_flush=True), and the tab updates while the task runs. See Reports. To log to an experiment tracker instead, see the Weights & Biases and MLflow integrations.

Hand the model to serving

Related

• LLM fine-tuning with LoRA and QLoRA: the full pipeline on this page, with live reports.
• Fine-tuning a VLM: multi-node fine-tuning with DeepSpeed and Lightning.
• BERT emotion classification: fine-tuning an encoder with the Hugging Face Trainer.
• RT-DETR object detection: fine-tuning a detector with cached data prep, live training charts, and mAP evaluation.
• Brain tumor MRI classification: two-phase transfer learning with checkpoints that survive preemption.
• Reinforcement learning: post-training with rewards instead of labels.