FLUX.2-klein-base-9B Layer Offloading seems broken

[StablePerception]

Thank you for adding Flux2 Klein support. However at least for the 9B Version Layer Offloading seems broken. My PC has a 16 GB Nvidia GPU and 78 GB of RAM.

I got Transformer and Text encoder Quantizations at 3 Bit, Layer offloading for the transformer on 100% and Cache Text embeddings enabled. But I still get a OOM during transformer loading and also the RAM (not VRAM) is not getting fuller when loading the transformer.

Hope you can fix this issue.

[jremb]

Seems that self.quantize_device is set to cuda, regardless of the low_vram setting in the config. So, when transformer.to(self.quantize_device, dtype=dtype) is called in the Flux2Model.load_model, it's trying to load the model onto the GPU. If self.quantize_device is set to 'cpu' (if you have low vram), it should work correctly.

[hugodopradofernandes]

Adding this to the line 82 of the file extensions_built_in/diffusion_models/flux2/flux2_model.py did the job, as mentioned in the previous comment, below the self.use_raw_control_images = True:

self.quantize_device = "cpu"

Section sample with the added line:

    super().__init__(
        device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs
    )
    self.is_flow_matching = True
    self.is_transformer = True
    self.target_lora_modules = ["Flux2"]
    # control images will come in as a list for encoding some things if true
    self.has_multiple_control_images = True
    # do not resize control images
    self.use_raw_control_images = True
    self.quantize_device = "cpu"

# static method to get the noise scheduler
@staticmethod
def get_train_scheduler():
    return CustomFlowMatchEulerDiscreteScheduler(**scheduler_config)

[hugodopradofernandes]

I had also to fix it in the text encoder loading that was also going to the GPU directly:
file: extensions_built_in/diffusion_models/flux2/flux2_klein_model.py
Add the lines:
line 36:
self.quantize_device = "cpu"

comment line 48 and add line 49:
text_encoder.to(self.quantize_device, dtype=dtype)

Code block will be like:

    # use the new format on this new model by default
    self.use_old_lokr_format = False
    self.quantize_device = "cpu"

def load_te(self):
    if self.flux2_klein_te_path is None:
        raise ValueError("flux2_klein_te_path must be set for Flux2KleinModel")
    dtype = self.torch_dtype
    self.print_and_status_update("Loading Qwen3")

    text_encoder: Qwen3ForCausalLM = Qwen3ForCausalLM.from_pretrained(
        self.flux2_klein_te_path,
        torch_dtype=dtype,
    )
    #text_encoder.to(self.device_torch, dtype=dtype)
    text_encoder.to(self.quantize_device, dtype=dtype)

    flush()

I belive this shoudn't be the final fix in the tool, as I'm loading it directly to CPU without any check.... it seems that the "if" block after that does the check, but after trying loading. Then it may require check the quant parameter before loading, and then load appropiratelly.

But it's working here, training running :)

[jremb]

Looking at how this is handled in other parts of the code, like ltx2.py, this may be more consistent.

Instead of changing behavior of self.quantize_device, add check for low_vram:

# Around line 158 in flux2_model.py
if not self.model_config.low_vram:
            transformer.to(self.quantize_device, dtype=dtype)

Same in Flux2KleinModel.load_te:

# Around line 47
if not self.model_config.low_vram:
            text_encoder.to(self.device_torch, dtype=dtype)

The rest of the code now properly handles quantization. And to more closely match the pattern in ltx2.py, can also add a flush() at flux2_model.py line 185 and, then, around line 238, after tokenizer = [pipe.tokenizer]:

if not self.low_vram: # Add this check
    pipe.transformer = pipe.transformer.to(self.device_torch) # Pull this up from underneath `text_encoder[0].eval()`

flush()
# just to make sure everything is on the right device and dtype
text_encoder[0].to(self.device_torch)
text_encoder[0].requires_grad_(False)
text_encoder[0].eval()

flush()

With the above changes I've successfully run training with 16GB VRAM and 64 RAM.

[valden80]

Just wanted to say 'Thank you!' to you both: @hugodopradofernandes, @jremb
It's really working, and with those I can train lora's for 9B Klein at the RTX 3060 12GB and 64GB RAM. It's still pretty quick, almost the same as for ZIT training.
I hope these fixes will be added to main branch soon...

[iamrohitanshu]

@valden80 Hello, I have the same hardware. Are you training t2i (single dataset), or i2i (target-control images dataset)?
Also, can you please share your config (what quantization and resolution) and what speeds are you getting?

[srjghseruikg]

This really needs to be merged, I can't train at all due to these OOM errors and i have more VRAM than @valden80 lol

[valden80]

@iamrohitanshu Rank 64, BS 4, 8bit transformer, 8bit encoder (100% offloading for both) - 10-17 s/it. Most time around 14. Single dataset for T2I, 512х512 image sizes. 8-9GB VRAM used. No OOM or other errors, process is stable.
I just trying it for now, so starting with simple case. There is still some free VRAM left to utilize...

[iamrohitanshu]

@valden80 Thanks! I'm able to use t2i as well. I tried 4-bit quantization, so didn't need to offload either (it might have happened automatically though).
Can you please try i2i? Because that's what causing problems for me. It doesn't go past the first iteration.

[valden80]

@iamrohitanshu I don't have dataset for it, and don't done anything other then t2i training before. I'm pretty much newbie. So I just don't know - how to do it in a right way.

[iamrohitanshu]

Oh okay. It's almost exactly the same, you just give different datasets for the target and control folders in the UI.
For reference, you can check out Ostris's youtube channel - specifically the video for Qwen Image Edit

My i2i training started once for 4b model, giving about 17 s/it speed. I then stopped it to disable the samples, and it just became so slow after re-starting it, like more than 100 s/it. So I just had to cancel it. I don't know where's the issue.

[valden80]

@iamrohitanshu If you are on Windows it's likely issue with automatic usage of RAM for offload when VRAM is overfilled at the driver level (can be disabled in nvidia control panel). I'm on Linux, and there will be just CUDA OOM in such situation, as there is no such feature in the drivers.

[iamrohitanshu]

That would make sense if my RAM was full, but the stats are exactly the same as when I was getting 17 s/it speed.
As you can see in the screenshot, my RAM has much empty space.

[valden80]

@iamrohitanshu No, it's not about your RAM. It's about your VRAM. And it's full, like I thought before (98.3%). On windows, in new nvidia drivers exists automatic offload - when VRAM is full, driver trying to offload some of it to RAM. It's extremely slow even if not fitting just a little of it, but CUDA OOM not happens in such case. But for AI tasks it's usually useless feature, as even without CUDA OOM it making it unusable slow - like your 100 sec/it, so disabling it may be a good thing to do. (You will be just getting CUDA OOM error instantly on overflows, without painful slowdowns and waiting for nothing.)

P.S. CUDA OOM - it's OOM for video card memory, not for normal RAM.

[iamrohitanshu]

@valden80 I'm sorry. I already knew that, it was a lapse in my reading comprehension.

The whole point of my post was that my training was giving 17s/it for about 150 steps, but when I paused and restarted it, it became very slow. The VRAM and RAM usage should have been the same, though I didn't check. It won't make sense that it'll increase for no reason. I wasn't using the GPU for anything else, and my monitor runs on a separate GPU.