What are good configs for training UNet3DConditionModel on 8 GB VRAM? (64x64x64 inputs)

[RandomGamingDev]

What are good configs for training UNet3DConditionModel on 8 GB VRAM? (64x64x64 inputs)

More specifically for this project I'm looking to use HuggingFace's UNet3DConditionModel on my home PC on a RTX 3060 Ti for voxel generation of small (64x64x64) models, but I'm struggling to find a good setup that would fit within the VRAM of my GPU. I doubt the model will be very effective with the space limitations I have so I'll probably swap to another model later, but I'd like to try out a setup with the UNet3DConditionModel first so any recommendations that don't require a more powerful computer?

[BenjaminBossan]

I think there is some confusion. PEFT is library that can help reduce the amount of memory when training models. Your question sounds like it aims at running the model. PEFT cannot help you with that.

I don't know what packages you use for running your model, whether you use one of the UIs or diffusers. In case of a UI like automatic1111, search their options for VRAM saving settings like --lowvram. If you use diffusers, check the diffusers docs or ask on the diffusers discussion board.

[RandomGamingDev]

I think there is some confusion. PEFT is library that can help reduce the amount of memory when training models. Your question sounds like it aims at running the model. PEFT cannot help you with that.

I don't know what packages you use for running your model, whether you use one of the UIs or diffusers. In case of a UI like automatic1111, search their options for VRAM saving settings like --lowvram. If you use diffusers, check the diffusers docs or ask on the diffusers discussion board.

Ah, sorry I phrased it wrongly. Yes, I'm looking to train the model, not just to run it.

[BenjaminBossan]

I see. 8GB of VRAM is indeed very little, even when using PEFT, but as with all things, it depends on there are many knobs to turn. The bigger issue, though, is that this model most likely uses modules that are not supported by any implemented PEFT method, e.g. DownBlock3D. If you could show the repr of your specific model, I could tell you if any PEFT method could actually be applied.

[RandomGamingDev]

While trying to get the model to run I minimized it down to this which barely runs on my CPU on 32 GB RAM with ~7-8 GB used up by my OS and activities:

model = UNet3DConditionModel(
    sample_size=config.image_size, # 64
    in_channels=4,
    out_channels=4,
    layers_per_block=1,
    block_out_channels=(16, 32),
    down_block_types=(
        "DownBlock3D",
        "DownBlock3D",
    ),
    up_block_types=(
        "UpBlock3D",
        "UpBlock3D",
    ),
    norm_num_groups=16,
    attention_head_dim=16
)

This is also being trained from scratch with no pretraining

[BenjaminBossan]

Okay, so I did a count on the number of parameters per layer type on this model:

• torch.nn.modules.conv.Conv2d: 140772,
• torch.nn.modules.linear.Linear: 461472,
• torch.nn.modules.normalization.GroupNorm: 2848,
• torch.nn.modules.normalization.LayerNorm: 1152,
• torch.nn.modules.conv.Conv3d: 71488

So most of the parameters are on Conv2d and Linear, not on Conv3d (which is not supported), so in theory, using LoRA could be helpful.

This is also being trained from scratch with no pretraining

This is a big problem. PEFT is intended for fine-tuning, i.e. taking a pretrained model and adapting it to your specific problem. You will almost certainly not succeed when training from scratch. I imagine you cannot fully train your model because of the memory constraints. I would suggest that you go on a search for a model that is pretrained on a similar domain as what you're interested in. You could check the HF forums or discord. If you find one and it needs further tuning, then using PEFT makes sense.

[RandomGamingDev]

This is a big problem. PEFT is intended for fine-tuning, i.e. taking a pretrained model and adapting it to your specific problem. You will almost certainly not succeed when training from scratch. I imagine you cannot fully train your model because of the memory constraints. I would suggest that you go on a search for a model that is pretrained on a similar domain as what you're interested in. You could check the HF forums or discord. If you find one and it needs further tuning, then using PEFT makes sense.

Yeah, PEFT isn't a realistic option here since I'd prefer training from scratch for this since I don't really see any pretrained models that fit my goals right now. I'll just swap to another model that I have in mind rn. Thanks.

(Also, I just realized that I posted on the wrong HuggingFace discussion page. I meant to post on diffusers and must've misclicked. Sorry about that.)

[BenjaminBossan]

No worries. Good luck with your task.