- [FLUX](#jump1) - [模型介绍](#模型介绍) - [微调](#微调) - [环境搭建](#环境搭建) - [微调](#jump2) - [性能](#性能) - [推理](#推理) - [环境搭建及运行](#环境搭建及运行) - [性能](#jump3) - [环境变量声明](#环境变量声明) - [引用](#引用) - [公网地址说明](#公网地址说明) # FLUX ## 模型介绍 [FLUX.1 dev](https://blackforestlabs.ai/announcing-black-forest-labs/) 是一种基于Rectified Flow Transformers (矫正流) 的生成模型。 - 参考实现： ```shell url=https://github.com/huggingface/diffusers commit_id= a98a839de75f1ad82d8d200c3bc2e4ff89929081 ``` ## 微调 ### 环境搭建 【模型开发时推荐使用配套的环境版本】 请参考[安装指南](https://gitee.com/ascend/MindSpeed-MM/blob/master/docs/user-guide/installation.md) 1. 软件与驱动安装 ```bash # python3.8 conda create -n test python=3.8 conda activate test # 安装 torch 和 torch_npu，注意要选择对应python版本、x86或arm的torch、torch_npu及apex包 pip install torch-2.1.0-cp38-cp38m-manylinux2014_aarch64.whl pip install torch_npu-2.1.0*-cp38-cp38m-linux_aarch64.whl # apex for Ascend 参考 https://gitee.com/ascend/apex # 建议从原仓编译安装 # 将shell脚本中的环境变量路径修改为真实路径，下面为参考路径 source /usr/local/Ascend/ascend-toolkit/set_env.sh ``` 2. 克隆仓库到本地服务器 ```shell git clone --branch 2.1.0 https://gitee.com/ascend/MindSpeed-MM.git ``` 3. 模型搭建 3.1 【下载 FLUX [GitHub参考实现](https://github.com/huggingface/diffusers) 在模型根目录下执行以下命令，安装模型对应PyTorch版本需要的依赖】 ```shell git clone https://github.com/huggingface/diffusers.git cd diffusers git checkout a98a839de75f1ad82d8d200c3bc2e4ff89929081 cp -r ../MindSpeed-MM/examples/diffusers/flux/* ./examples/dreambooth ``` 【主要代码路径】 ```shell code_path=examples/dreambooth/ ``` 3.2【安装其余依赖库】 ```shell pip install -e . vim examples/dreambooth/requirements_flux.txt #修改版本：torchvision==0.16.0, torch==2.1.0, accelerate==0.33.0, 添加deepspeed==0.15.2 pip install -r examples/dreambooth/requirements_flux.txt # 安装对应依赖 ``` ## 微调 1. 【准备微调数据集】 - 用户需自行获取并解压[pokemon-blip-captions](https://huggingface.co/datasets/lambdalabs/pokemon-blip-captions/tree/main)数据集，并在以下启动shell脚本中将`dataset_name`参数设置为本地数据集的绝对路径 ```shell dataset_name="pokemon-blip-captions" # 数据集 路径 ``` - pokemon-blip-captions数据集格式如下: ```shell pokemon-blip-captions ├── dataset_infos.json ├── README.MD └── data └── train-001.parquet ``` - 只包含图片的训练数据集，如deepspeed脚本使用训练数据集dog:[下载地址](https://huggingface.co/datasets/diffusers/dog-example)，并将dog文件夹转移到`examples/dreambooth/`目录下 ```shell input_dir="dog" # 数据集路径 ``` ```shell dog ├── alvan-nee-*****.jpeg ├── alvan-nee-*****.jpeg ``` > **说明：** >该数据集的训练过程脚本只作为一种参考示例。 > - 如用自己的微调数据集，需在shell脚本中将`input_dir`修改为`dataset_name`： ```shell dataset_name="/path/customized_datasets" # 数据集路径 ``` 在shell脚本`accelerate launch`目录下（70行左右）将`instance_data_dir=$instance_dir \` 修改为 `dataset_name=$dataset_name`，并将`instance_prompt`与`validation_prompt`改为与自己数据集所匹配的prompt: ```shell # Example accelerate launch --config_file ${config_file} \ ./train_dreambooth_flux.py \ --pretrained_model_name_or_path=$model_name \ --dataset_name=$dataset_name \ --instance_prompt="a prompt that is suitable for your own dataset" \ --validation_prompt="a validation prompt based on your own dataset" \ ``` 2. 【配置 FLUX 微调脚本】 联网情况下，微调模型可通过以下步骤下载。无网络时，用户可访问huggingface官网自行下载[FLUX.1-dev模型](https://huggingface.co/black-forest-labs/FLUX.1-dev) `model_name`模型 ```bash export model_name="black-forest-labs/FLUX.1-dev" # 预训练模型路径 ``` 获取对应的微调模型后，在以下shell启动微调脚本中将`model_name`参数设置为本地预训练模型绝对路径 ```shell model_name="black-forest-labs/FLUX.1-dev" # 预训练模型路径 batch_size=16 max_train_steps=5000 mixed_precision="bf16" # 混精 resolution=256 config_file="pretrain_${mixed_precision}_accelerate_config.yaml" # accelerate launch --config_file ${config_file} \ 目录下 --dataloader_num_workers=0 \ # 请基于系统配置与数据大小进行调整num workers ``` 3. 【修改代码文件】 1. 在 `src/diffusers/models/embeddings.py` 文件里，在 `class FluxPosEmbed(nn.Module):` 下的 **第813行左右** 找到代码： `freqs_dtype = torch.float32 if is_mps else torch.float64` 进行修改, 请修改为：`freqs_dtype = torch.float32` ```shell # 修改为freqs_dtype = torch.float32 vim src/diffusers/models/embeddings.py ``` ```python freqs_dtype = torch.float32 # 813行附近 # freqs_dtype = torch.float32 if is_mps else torch.float64 # 原代码 ``` 2. 打开`train_dreambooth_flux.py`或`train_dreambooth_lora_flux_advanced.py`文件 ```shell cd examples/dreambooth/ # 从diffusers目录进入dreambooth目录 vim train_dreambooth_flux.py # 进入Python文件 # 如是flux lora，需先进入advanced_diffusion_training目录 vim ../advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py # 进入Python文件 ``` - 在import栏`if is_wandb_available():`上方（62行附近添加代码） ```python # 添加代码到train_dreambooth_flux.py 62行附近 from patch_flux import TorchPatcher, config_gc, create_save_model_hook TorchPatcher.apply_patch() config_gc() if is_wandb_available(): # 原代码 import wandb ``` - 在log_validation里修改`pipeline = pipeline.to(accelerator.device)`，`train_dreambooth_flux.py`在171行附近 ```python # 修改pipeline为： pipeline = pipeline.to(accelerator.device, dtype=torch_dtype) # pipeline = pipeline.to(accelerator.device) # 原代码 ``` 3. 【Optional】Ubuntu系统需在1701行附近 添加 `accelerator.print("")` ```python if global_step >= args.max_train_steps: # 原代码 break accelerator.print("") # 添加 ``` 4. 【Optional】模型checkpoint saving保存 【因模型较大 如不需要`checkpointing_steps`，请设置其大于`max_train_steps`, 避免开启】 ```shell --checkpointing_steps=50000 \ # 修改50000步为所需要步数 ``` 【如需保存checkpointing请修改代码】 ```shell vim train_dreambooth_flux.py #（1669行附近） vim ../advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py #（2322行附近） ``` - 在文件上方的import栏增加`DistributedType`在`from accelerate import Acceleratore`后 （30行附近） - 在`if accelerator.is_main_process`后增加 `or accelerator.distributed_type == DistributedType.DEEPSPEED` (1669/2322行附近)，并在`if args.checkpoints_total_limit is not None`后增加`and accelerator.is_main_process` ```python from accelerate import Accelerator, DistributedType # from accelerate import Accelerator # 原代码 if accelerator.is_main_process or accelerator.distributed_type == DistributedType.DEEPSPEED: # if accelerator.is_main_process: # 原代码 if global_step % args.checkpointing_steps == 0: # 原代码 不进行修改 if args.checkpoints_total_limit is not None and accelerator.is_main_process: # 添加 ``` Lora任务需调用patch任务进行权重保存： 在`train_dreambooth_lora_flux_advanced.py`文件中找到代码`accelerator.register_save_state_pre_hook(save_model_hook)`进行修改(1712行附近)，复制粘贴以下代码： ```python # 添加 save_Model_Hook = create_save_model_hook( accelerator=accelerator, unwrap_model=unwrap_model, transformer=transformer, text_encoder_one=text_encoder_one, args=args, weight_dtype=weight_dtype ) accelerator.register_save_state_pre_hook(save_Model_Hook) # 修改 # accelerator.register_save_state_pre_hook(save_model_hook) # 原代码 accelerator.register_load_state_pre_hook(load_model_hook) # 原代码 不修改 ``` 更改shell脚本： ```shell export HCCL_CONNECT_TIMEOUT=1200 # 大幅调高HCCL_CONNECT_TIMEOUT (如5000) export HCCL_EXEC_TIMEOUT=17000 --checkpointing_steps=50000 \ # 修改50000步为所需要步数 ``` 4. 【启动 FLUX 微调脚本】 本任务主要提供flux_dreambooth与flux_dreambooth_lora微调脚本，支持多卡训练。 启动FLUX dreambooth微调脚本 ```shell bash finetune_flux_dreambooth_deepspeed_bf16.sh ``` 启动FLUX dreambooth_lora微调脚本 ```shell bash finetune_flux_dreambooth_lora_deepspeed_bf16.sh ``` ### 性能 #### 吞吐 FLUX 在 **昇腾芯片** 和 **参考芯片** 上的性能对比： | 芯片 | 卡数 | 任务 | FPS | batch_size | AMP_Type | Torch_Version | deepspeed | |:---:|:---:|:----------:|:-----:|:----------:|:---:|:---:|:---:| | Atlas 900 A2 PODc | 8p | Flux-全参微调 | 55.23 | 16 | bf16 | 2.1 | ✔ | | 竞品A | 8p | Flux-全参微调 | 53.65 | 16 | bf16 | 2.1 | ✔ | ## 推理 ### 环境搭建及运行 **同微调对应章节** ```shell cd examples/dreambooth/ # 从diffusers目录进入dreambooth目录 ``` 【FLUX模型推理】 ```shell vim infer_flux_text2img_bf16.py # 进入运行推理的Python文件 ``` 1. 修改路径 ```python MODEL_PATH = "/black-forest-labs/FLUX.1-dev" # FLUX模型路径 ``` 2. 运行代码 ```shell python infer_flux_text2img_bf16.py ``` 【DREAMBOOTH微调FLUX模型推理】 ```shell vim infer_flux_text2img_dreambooth_bf16.py ``` 1. 修改路径 ```python MODEL_PATH = "./output_FLUX_dreambooth" # Dreambooth微调保存模型路径 ``` 2. 运行代码 ```shell python infer_flux_text2img_dreambooth_bf16.py ``` 【lora微调FLUX模型推理】 ```shell vim infer_flux_text2img_lora_bf16.py ``` 1. 修改路径 ```python MODEL_PATH = "./FLUX" # Flux 模型路径 LORA_WEIGHTS = "./output/pytorch_lora_weights.safetensors" # LoRA权重路径 ``` 2. 运行代码 ```shell python infer_flux_text2img_lora_bf16.py ``` 【分布式推理】 ```shell vim infer_flux_text2img_distrib.py ``` - 修改模型权重路径 model_path为模型权重路径或微调后的权重路径 - 如lora微调 可将lora_weights修改为Lora权重路径 ```python model_path = "/black-forest-labs/FLUX.1-dev" # 模型权重/微调权重路径 lora_weights = "/pytorch_lora_weights.safetensors" # Lora权重路径 ``` - 启动分布式推理脚本 - 因使用accelerate进行分布式推理，config可设置：`--num_processes=卡数`，`num_machines=机器数`等 ```shell accelerate launch --num_processes=4 infer_flux_text2img_distrib.py # 单机四卡进行分布式推理 ``` ### 性能 | 芯片 | 卡数 | 任务 | E2E（it/s） | AMP_Type | Torch_Version | |:---:|:---:|:----------:|:-----:|:---:|:---:| | Atlas 900 A2 PODc |8p | 文生图 | 1.16 | bf16 | 2.1 | | 竞品A | 8p | 文生图 | 1.82 | bf16 | 2.1 | | Atlas 900 A2 PODc |8p | 文生图微调 | 1.12 | bf16 | 2.1 | | 竞品A | 8p | 文生图微调 | 1.82 | bf16 | 2.1 | ## 环境变量声明 ASCEND_SLOG_PRINT_TO_STDOUT： 是否开启日志打印， 0：关闭日志打屏，1：开启日志打屏 ASCEND_GLOBAL_LOG_LEVEL： 设置应用类日志的日志级别及各模块日志级别，仅支持调试日志。0：对应DEBUG级别，1：对应INFO级别，2：对应WARNING级别，3：对应ERROR级别，4：对应NULL级别，不输出日志 TASK_QUEUE_ENABLE： 用于控制开启task_queue算子下发队列优化的等级，0：关闭，1：开启Level 1优化，2：开启Level 2优化 COMBINED_ENABLE： 设置combined标志。设置为0表示关闭此功能；设置为1表示开启，用于优化非连续两个算子组合类场景 HCCL_WHITELIST_DISABLE： 配置在使用HCCL时是否开启通信白名单，0：开启白名单，1：关闭白名单 CPU_AFFINITY_CONF： 控制CPU端算子任务的处理器亲和性，即设定任务绑核，设置0或未设置：表示不启用绑核功能， 1：表示开启粗粒度绑核， 2：表示开启细粒度绑核 HCCL_CONNECT_TIMEOUT: 用于限制不同设备之间socket建链过程的超时等待时间，需要配置为整数，取值范围[120,7200]，默认值为120，单位s ACLNN_CACHE_LIMIT： 配置单算子执行API在Host侧缓存的算子信息条目个数 TOKENIZERS_PARALLELISM： 用于控制Hugging Face的transformers库中的分词器（tokenizer）在多线程环境下的行为 PYTORCH_NPU_ALLOC_CONF： 控制缓存分配器行为 OMP_NUM_THREADS： 设置执行期间使用的线程数 ## 引用 ### 公网地址说明 代码涉及公网地址参考 [公网地址](https://gitee.com/ascend/MindSpeed-MM/blob/master/docs/public_address_statement.md)