Which DeepSpeed level is necessary to fine-tune flux dev on 1XH100 90GB VRAM?

[vedantroy]

I'm trying to fine-tune flux dev on a single H100. Is this the correct deepspeed config?

Compute environment | This machine
Machine type | No distributed training
Run on CPU only | NO
Use torch dynamo | yes
Dynamo backend | inductor
Customize torch.compile defaults | no
Use DeepSpeed | yes
Specify DeepSpeed config JSON | no
DeepSpeed ZeRO stage | 3
Optimizer states offload | cpu
Parameters offload | cpu
Gradient accumulation steps | 1
Use gradient clipping | no
Save 16-bit model weights | no
Enable deepspeed.zero.Init | no
Enable MoE training | no
Number of GPUs | 1
Mixed precision | bf16

Additionally, my training is getting stuck here (after "configured dataset"). What should I do? gpustat -cupi indicates that GPU util is at 0%. Not sure if it matters, but I'm using a machine with 128GB of RAM.

You set `add_prefix_space`. The tokenizer needs to be converted from the slow tokenizers
2025-02-13 12:08:35,225 [INFO] Loading OpenAI CLIP-L text encoder from /root/comfy/models/training/text_encoder..
2025-02-13 12:08:35,250 [INFO] Loading T5 XXL v1.1 text encoder from /root/comfy/models/training/text_encoder_2..
Loading checkpoint shards: 100%|__________________________________________________________________________________________________________________________________________| 2/2 [00:00<00:00,  4.38it/s]
2025-02-13 12:08:38,109 [INFO] Moving text encoder to GPU.
2025-02-13 12:08:38,159 [INFO] Moving text encoder 2 to GPU.
2025-02-13 12:08:40,114 [INFO] Loading data backend config from /finetune/processed_training_v0_1739448486/datasetbackend.json
2025-02-13 12:08:40,115 [INFO] Configuring text embed backend: text-embed-cache
Loading pipeline components...: 100%|____________________________________________________________________________________________________________________________________| 5/5 [00:00<00:00, 755.49it/s]
2025-02-13 12:08:40,125 [INFO] (Rank: 0) (id=text-embed-cache) Listing all text embed cache entries
2025-02-13 12:08:40,409 [INFO] Pre-computing null embedding
2025-02-13 12:08:46,163 [INFO] Completed loading text embed services.
2025-02-13 12:08:46,164 [INFO] Configuring data backend: dataset-v0_1739448486
2025-02-13 12:08:46,164 [INFO] (id=dataset-v0_1739448486) Loading bucket manager.
2025-02-13 12:08:46,165 [WARNING] No cache file found, creating new one.
2025-02-13 12:08:46,166 [INFO] (id=dataset-v0_1739448486) Refreshing aspect buckets on main process.
2025-02-13 12:08:46,166 [INFO] Discovering new files...
2025-02-13 12:08:46,198 [INFO] Compressed 0 existing files from 0.
2025-02-13 12:08:47,067 [INFO] Image processing statistics: {'total_processed': 232, 'skipped': {'already_exists': 0, 'metadata_missing': 0, 'not_found': 0, 'too_small': 0, 'other': 0}}
2025-02-13 12:08:47,070 [INFO] Enforcing minimum image size of 1.048576. This could take a while for very-large datasets.
2025-02-13 12:08:47,072 [INFO] Completed aspect bucket update.
2025-02-13 12:08:47,075 [INFO] Configured backend: {'id': 'dataset-v0_1739448486', 'config': {'repeats': 500, 'crop': False, 'crop_aspect': 'square', 'crop_style': 'random', 'disable_validation': False, 'resolution': 1.048576, 'resolution_type': 'area', 'caption_strategy': 'textfile', 'instance_data_dir': '/finetune/processed', 'maximum_image_size': None, 'target_downsample_size': None, 'config_version': 2}, 'dataset_type': 'image', 'data_backend': <helpers.data_backend.local.LocalDataBackend object at 0x78691b71ca00>, 'instance_data_dir': '/finetune/processed', 'metadata_backend': <helpers.metadata.backends.discovery.DiscoveryMetadataBackend object at 0x7869081c7100>}
(Rank: 0)  | Bucket     | Image Count (per-GPU)
------------------------------
(Rank: 0)  | 1.0        | 232
2025-02-13 12:08:47,077 [INFO] (id=dataset-v0_1739448486) Collecting captions.
2025-02-13 12:08:47,094 [INFO] (id=dataset-v0_1739448486) Initialise text embed pre-computation using the textfile caption strategy. We have 232 captions to process.
2025-02-13 12:08:56,432 [INFO] (id=dataset-v0_1739448486) Completed processing 232 captions.
2025-02-13 12:08:56,432 [INFO] (id=dataset-v0_1739448486) Creating VAE latent cache.
2025-02-13 12:08:56,436 [INFO] (id=dataset-v0_1739448486) Discovering cache objects..
(id=dataset-v0_1739448486) Bucket 1.0 caching results: {'not_local': 0, 'already_cached': 0, 'cached': 4, 'total': 232}
2025-02-13 12:09:09,786 [INFO] Configured backend: {'id': 'dataset-v0_1739448486', 'config': {'repeats': 500, 'crop': False, 'crop_aspect': 'square', 'crop_style': 'random', 'disable_validation': False, 'resolution': 1.048576, 'resolution_type': 'area', 'caption_strategy': 'textfile', 'instance_data_dir': '/finetune/processed', 'maximum_image_size': None, 'target_downsample_size': None, 'config_version': 2, 'hash_filenames': True}, 'dataset_type': 'image', 'data_backend': <helpers.data_backend.local.LocalDataBackend object at 0x78691b71ca00>, 'instance_data_dir': '/finetune/processed', 'metadata_backend': <helpers.metadata.backends.discovery.DiscoveryMetadataBackend object at 0x7869081c7100>, 'train_dataset': <helpers.multiaspect.dataset.MultiAspectDataset object at 0x7869081c7460>, 'sampler': <helpers.multiaspect.sampler.MultiAspectSampler object at 0x7869081c7640>, 'train_dataloader': <torch.utils.data.dataloader.DataLoader object at 0x78691b790100>, 'text_embed_cache': <helpers.caching.text_embeds.TextEmbeddingCache object at 0x78691e91f4c0>, 'vaecache': <helpers.caching.vae.VAECache object at 0x78691ad1dde0>}

[bghira]

probably inductor makes it get stuck i would guess

[vedantroy]

Yep, that prevented the stall. For full fine-tuning, do you have recommendations on the machine / accelerate or deepspeed config?

E.g

• How many GPUs (I'm using H100s)
• Should we use DeepSpeed? No Deepspeed? If we don't use DeepSpeed, should we use FSDP? etc.How many different machines will you use (use more than 1 for multi-node training)

[bghira]

i think 2x8 H100 is the minimum for really cranking a full rank tune and getting meaningfully high quality results.

FluxBooru I trained on at first, 8xH100 and then 16xH100 and the difference was so noticeable that I would have rather trained on 64xH100 if only I'd had them.