Skip to content

v3.0.2.md

Latest commit

 

History

History
90 lines (66 loc) · 2.3 KB

v3.0.2.md

File metadata and controls

90 lines (66 loc) · 2.3 KB
Raw

v3.0.2 Changelog

Release Date: 2025-10-26

This release primarily fixes FSDP training issues with LoRA checkpoint saving, while also documenting performance benchmarks for different training strategies.

Bug Fixes

  • Fixed FSDP LoRA checkpoint saving issues that prevented proper model weights storage
  • Resolved compatibility problems between FSDP sharding and LoRA adapters

New Features

  • Added support for FP4 DDP training configuration
  • Enhanced FSDP checkpoint management for LoRA weights

Improvements

  • Updated FSDP memory optimization plan with FP4 DDP benchmark results
  • Added performance metrics to TODO tracking

Documentation

  • Enhanced FSDP memory optimization documentation with FP4 DDP benchmarks
  • Updated TODO.md with performance metrics tracking item

Technical Details

FSDP LoRA Checkpoint Saving

  • Implemented custom state consolidation for FSDP-trained LoRA adapters
  • Added safeguards to ensure proper weight gathering before checkpoint saving
  • Fixed tensor sharding issues during LoRA adapter state extraction

Performance Benchmarks

  • BF16 DDP training configuration:

    • Memory usage: >24GB per GPU (causes OOM on RTX 4090)
    • Not feasible on RTX 4090 hardware

  • FP4 DDP training configuration:

    • Memory usage: ~25GB per GPU
    • Throughput: ~0.4 FPS
    • Suitable for environments where FSDP implementation is challenging

  • BF16 FSDP training configuration:

    • Memory usage: ~10GB per GPU
    • Throughput: ~1.7 FPS
    • Best balance of memory efficiency and performance

Configuration Example

# FP4 DDP 配置
NCCL_P2P_DISABLE=1 NCCL_IB_DISABLE=1 \
CUDA_VISIBLE_DEVICES=0,1 \
accelerate launch \
  --num_processes 2 \
  --mixed_precision bf16 \
  -m qflux.main --config $config_file
# FP4 DDP 模型配置
model:
  pretrained_model_name_or_path: lrzjason/flux-kontext-nf4
  quantize: false
# BF16 FSDP 配置（修复了LoRA checkpoint保存）
CUDA_VISIBLE_DEVICES=0,1 \
accelerate launch \
  --num_processes 2 \
  --mixed_precision bf16 \
  --use_fsdp \
  --fsdp_sharding_strategy 1 \
  -m qflux.main --config $config_file
# BF16 FSDP 模型配置
model:
  pretrained_model_name_or_path: black-forest-labs/FLUX.1-Kontext-dev
  quantize: false

Contributors

  • Team members involved in performance testing and documentation