Lumina cache checkpointing

[rockerBOO]

• Add cache using strategy and cache to disk
• Gradient checkpointing
• Gemma offloading
• Refactor of Lumina forward to be easier to read
• FP8 support

I noticed packing latents didn't work and the code didn't seem to reflect it being packed. Left it in the code but commented it out.

This PR implies #12 should go first as it builds off of that branch.

Related: kohya-ss#1927

[rockerBOO]

From this the model is training and I can load it in comfy and it seems "work" but the result is pretty noisy.

The loss when training starts at 2.5ish.

I think it's some variable of the noise and time shift.

timestep_sampling = "flux_shift" # and shift
discrete_flow_shift = 3.1582
model_prediction_type = "raw"
guidance_scale = 1.0

Tried these but might be something else needed here.

[rockerBOO]

not setting any of those values seems to do a little better. Also using huber loss with huber_c = 0.2 and huber_scale = 1.8. Trying some others but im mostly just experimenting without really knowing.

[rockerBOO]

without huber loss (using l2 loss) I think maybe something with the sampler too but a little closer

But the loss spiked quite a bit

[rockerBOO]

I tried the packed latents again. Seems it goes through if it packs and not unpack but the results were not good.