If you’ve been trying to train a custom FLUX.1-dev LoRA on a cloud GPU like RunPod, you’ve probably realized two things: the results are insanely good, but setting up the environment can feel like walking through a minefield of VRAM bottlenecks and cryptic YAML errors.

I spent the better part of a weekend tweaking configs on an NVIDIA A40 (48GB VRAM) instance so you don’t have to. Whether you hit a ValueError: Unknown process type right out of the gate or ran straight into a CUDA Out of Memory (OOM) wall mid-process, here is the straightforward, battle-tested guide to getting your LoRA trained safely and smoothly.

The Setup: What We’re Building

Before jumping into the commands, here is the hardware and dataset footprint used for this build:

• GPU: 1x NVIDIA A40 (48GB VRAM)
• Dataset: ~130 high-resolution images with matching .txt caption files
• Trigger Word:mystic_vibe (replace this with your custom trigger)
• Goal: A clean .safetensors LoRA file ready for ComfyUI or Forge.

Step 1: Environment & Dependency Quick-Fixes

First, clone the official ai-toolkit

repository and set up your Python virtual environment inside your RunPod workspace:

Bash

git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit

python3 -m venv venv
source venv/bin/activate

pip install --upgrade pip
pip install -r requirements.txt

Common Dependency Pitfall

Depending on your RunPod PyTorch base image, running python run.py might immediately throw a ModuleNotFoundError: No module named 'torchaudio'.

AI-Toolkit relies on torchaudio for dataset preprocessing pipelines. Fix this right away by installing it alongside the proper CUDA 12 wheel:

Bash

pip install torchaudio
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

Step 2: Fixing the “Unknown process type” Error

If you pull old configuration templates off forum posts, you might see an error like this when launching your job:

Plaintext

ValueError: config file is invalid. Unknown process type: train_network

AI-Toolkit uses sd_trainer for Diffusion/FLUX training jobs. Your top-level job key must be set to extension or train, and the process item type must be sd_trainer.

Make sure your dataset folder path is absolute (e.g., /workspace/ai-toolkit/dataset/your_folder).

Step 3: Preventing CUDA Out Of Memory (OOM) on 48GB VRAM

FLUX.1-dev is a massive model. Even with 48GB VRAM on an A40, loading the full FP16 Transformer alongside T5-XXL and CLIP text encoders can instantly spike memory usage past 44GB during initialization, causing a crash before epoch 1 even starts.

To fix this without destroying output quality, apply three critical tweaks in your YAML config:

1. quantize: true: Loads the FLUX transformer in 8-bit mode.
2. low_vram: true: Automatically offloads text encoders to system RAM after caching latents.
3. batch_size: 1 with gradient_accumulation_steps: 2: Keeps your VRAM footprint around ~24GB while maintaining an effective batch size of 2.

Also, before launching your script in the terminal, set PyTorch’s memory allocation flag to reduce fragmentation:

Bash

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

Step 4: The Complete, Working YAML Config

Create a file named config/mystic_flux.yaml inside your ai-toolkit folder and paste the following configuration. This includes automated image sample generation every time a checkpoint is saved.

YAML

job: extension
config:
  name: "mystic_flux_lora"
  training_folder: "output"
  device: "cuda:0"
  trigger_word: "mystic_vibe"
  process:
    - type: sd_trainer
      training_folder: "output"
      device: "cuda:0"
      trigger_word: "mystic_vibe"
      network:
        type: "lora"
        linear: 16
        linear_alpha: 16
      save:
        dtype: "bfloat16"
        save_every: 250
        max_step_saves_to_keep: 4
        push_to_hub: false
      
      # Sample generation during training
      sample:
        sampler: "flowmatch"
        sample_every: 250
        width: 1024
        height: 1024
        prompts:
          - "a photo of mystic_vibe standing in a neon-lit city street at night"
          - "a close up portrait of mystic_vibe, cinematic lighting, highly detailed"
        neg: ""
        seed: 42
        walk_seed: true
        guidance_scale: 3.5
        sample_steps: 20

      datasets:
        - folder_path: "/workspace/ai-toolkit/dataset/mystic_data"
          caption_ext: "txt"
          caption_dropout_rate: 0.05
          shuffle_tokens: false
          cache_latents_to_disk: true
          resolution:
            - 512
            - 768
            - 1024
      train:
        batch_size: 1
        gradient_accumulation_steps: 2
        steps: 1500
        train_unet: true
        train_text_encoder: false
        gradient_checkpointing: true
        noise_scheduler: "flowmatch"
        optimizer: "adamw8bit"
        lr: 0.0001
        dtype: "bf16"
      model:
        name_or_path: "black-forest-labs/FLUX.1-dev"
        is_flux: true
        quantize: true
        low_vram: true
      lr_scheduler:
        type: "cosine"
        warmup_steps: 100
      logger:
        type: "console"
        storage_location: "output/log"
meta:
  name: "mystic_flux_lora"
  version: '1.0'

Step 5: Launching the Training Job

With your dataset uploaded to /workspace/ai-toolkit/dataset/mystic_data and your config saved, run the training command:

Bash

python run.py config/mystic_flux.yaml

What to expect during execution:

1. Latent Caching: AI-Toolkit will read your 130+ images, encode them via VAE, and save cache files to disk.
2. Text Offloading: Thanks to low_vram: true, you’ll see system RAM usage rise briefly while VRAM drops back to a safe ~24GB.
3. Training & Sampling: Every 250 steps, the trainer will save a .safetensors file inside output/mystic_flux_lora/ and generate test images inside output/mystic_flux_lora/samples/.

Recap & Next Steps

Training FLUX LoRAs doesn’t require a cluster of enterprise GPUs if you dial in your memory settings correctly. By enabling 8-bit quantization, low VRAM offloading, and gradient accumulation, an A40 handles FLUX.1-dev smoothly with zero OOM crashes.

Once your 1,500 steps finish:

• Check the samples/ directory to evaluate which step interval (e.g., step 750 vs step 1250) captured your dataset best without overtraining.
• Download your chosen .safetensors checkpoint from the output/ folder and test it in your local ComfyUI workflow using your custom trigger word!