已合并
feat:add A5 script for flux DanceGRPO #2904
xubin创建于 23 天前
feat:add A5 script for flux DanceGRPO #2904
已合并
共 6 个文件变更+236-4
| @@ -199,6 +199,16 @@ bash examples/dancegrpo/posttrain_flux_dancegrpo.sh | |||
| 199 | 199 | ||
| 200 | 训练完成后，会在logs目录中生成运行日志文件，生成训练reward记录文件。 | 200 | 训练完成后，会在logs目录中生成运行日志文件，生成训练reward记录文件。 |
| 201 | 201 | ||
| 202 | +#### Atlas 950 服务器性能测试脚本 | ||
| 203 | + | ||
| 204 | +如需在 Atlas 950 服务器上进行纯性能测试，可使用 [posttrain_flux_dancegrpo_A5.sh](./posttrain_flux_dancegrpo_A5.sh) 脚本。该脚本默认通过 `--no-save_images` 关闭采样图片的落盘，以减少 I/O 开销、专注于性能评估： | ||
| 205 | + | ||
| 206 | +```bash | ||
| 207 | +bash examples/dancegrpo/posttrain_flux_dancegrpo_A5.sh | ||
| 208 | +``` | ||
| 209 | + | ||
| 210 | +> *注意：如需保留采样图片用于可视化，请将脚本中的 `--no-save_images \` 一行删除，恢复默认的图片保存行为（图片会写入 `./images/flux_grpo/step{N}/` 目录下）。* | ||
 | |||
| 211 | + | ||
| 202 | --- | 212 | --- |
| 203 | 213 | ||
| 204 | ## 性能数据 | 214 | ## 性能数据 |
| @@ -0,0 +1,117 @@ | |||
| 1 | +#!/bin/bash | ||
| 2 | + | ||
| 3 | +# 根据实际情况修改 ascend-toolkit 路径 | ||
| 4 | +source /usr/local/Ascend/cann/set_env.sh | ||
| 5 | +# 该变量只用于规避megatron对其校验，对npu无效 | ||
| 6 | +export CUDA_DEVICE_MAX_CONNECTIONS=1 | ||
| 7 | +export ASCEND_SLOG_PRINT_TO_STDOUT=0 | ||
| 8 | +export ASCEND_GLOBAL_LOG_LEVEL=3 | ||
| 9 | +export TASK_QUEUE_ENABLE=2 | ||
| 10 | +export COMBINED_ENABLE=1 | ||
| 11 | +export CPU_AFFINITY_CONF=2 | ||
| 12 | +export HCCL_CONNECT_TIMEOUT=1200 | ||
| 13 | +export NPU_ASD_ENABLE=0 | ||
| 14 | +export ASCEND_LAUNCH_BLOCKING=0 | ||
| 15 | +export ACLNN_CACHE_LIMIT=100000 | ||
| 16 | +export MULTI_STREAM_MEMORY_REUSE=2 | ||
| 17 | +export PYTORCH_NPU_ALLOC_CONF="expandable_segments:True" | ||
| 18 | +export HCCL_BUFFSIZE=800 | ||
| 19 | +export HCCL_OP_EXPANSION_MODE="CCU_SCHED" | ||
| 20 | + | ||
| 21 | +# 通过此配置选择使用的NPU卡,卡数需要与NPUS_PER_NODE相对应 | ||
| 22 | +export ASCEND_RT_VISIBLE_DEVICES="0,1,2,3,4,5,6,7" | ||
| 23 | +NPUS_PER_NODE=8 | ||
| 24 | +MASTER_ADDR=localhost | ||
| 25 | +MASTER_PORT=6000 | ||
| 26 | +NNODES=1 | ||
| 27 | +NODE_RANK=0 | ||
| 28 | +WORLD_SIZE=$(($NPUS_PER_NODE*$NNODES)) | ||
| 29 | + | ||
| 30 | + | ||
| 31 | +MM_DATA="./examples/dancegrpo/data_dancegrpo.json" | ||
| 32 | +MM_MODEL="./examples/dancegrpo/model_dancegrpo.json" | ||
| 33 | +MM_TOOL="./mindspeed_mm/tools/tools.json" | ||
| 34 | +LOAD_PATH="ckpt/flux" | ||
| 35 | +HPS_REWARD_SAVE_PATH="./hps_reward.txt" | ||
| 36 | + | ||
| 37 | +DISTRIBUTED_ARGS=" | ||
| 38 | + --nproc_per_node $NPUS_PER_NODE \ | ||
| 39 | + --nnodes $NNODES \ | ||
| 40 | + --node_rank $NODE_RANK \ | ||
| 41 | + --master_addr $MASTER_ADDR \ | ||
| 42 | + --master_port $MASTER_PORT | ||
| 43 | +" | ||
| 44 | + | ||
| 45 | +GPT_ARGS=" | ||
| 46 | + --seed 42 \ | ||
| 47 | + --load $LOAD_PATH \ | ||
| 48 | + --lr 1.0e-5 \ | ||
| 49 | + --train-iters 300 \ | ||
| 50 | + --weight-decay 0.0001 \ | ||
| 51 | +" | ||
| 52 | + | ||
| 53 | +MM_ARGS=" | ||
| 54 | + --mm-data $MM_DATA \ | ||
| 55 | + --mm-model $MM_MODEL \ | ||
| 56 | + --mm-tool $MM_TOOL | ||
| 57 | +" | ||
| 58 | + | ||
| 59 | +OUTPUT_ARGS=" | ||
| 60 | + --log-interval 1 \ | ||
| 61 | + --save-interval 10000 \ | ||
| 62 | + --eval-interval 10000 \ | ||
| 63 | + --eval-iters 5000 \ | ||
| 64 | + --ckpt-format torch \ | ||
| 65 | +" | ||
| 66 | + | ||
| 67 | +GRPO_ARGS=" | ||
| 68 | + --cache_dir data/.cache \ | ||
| 69 | + --gradient_checkpointing \ | ||
| 70 | + --train_batch_size 1 \ | ||
| 71 | + --num_latent_t 1 \ | ||
| 72 | + --sp_size 1 \ | ||
| 73 | + --train_sp_batch_size 1 \ | ||
| 74 | + --dataloader_num_workers 4 \ | ||
| 75 | + --gradient_accumulation_steps 4 \ | ||
| 76 | + --mixed_precision bf16 \ | ||
| 77 | + --cfg 0.0 \ | ||
| 78 | + --h 720 \ | ||
| 79 | + --w 720 \ | ||
| 80 | + --t 1 \ | ||
| 81 | + --sampling_steps 16 \ | ||
| 82 | + --eta 0.3 \ | ||
| 83 | + --lr_warmup_steps 0 \ | ||
| 84 | + --sampler_seed 1223627 \ | ||
| 85 | + --max_grad_norm 1.0 \ | ||
| 86 | + --use_hpsv2 \ | ||
| 87 | + --num_generations 12 \ | ||
| 88 | + --shift 3 \ | ||
| 89 | + --use_group \ | ||
| 90 | + --ignore_last \ | ||
| 91 | + --timestep_fraction 0.6 \ | ||
| 92 | + --init_same_noise \ | ||
| 93 | + --clip_range 1e-4 \ | ||
| 94 | + --adv_clip_max 5.0 \ | ||
| 95 | + --hps_reward_save $HPS_REWARD_SAVE_PATH \ | ||
| 96 | + --sample_batch_size 12 \ | ||
| 97 | + --no-save_images \ | ||
| 98 | +" | ||
| 99 | + | ||
| 100 | +logfile=$(date +%Y%m%d)_$(date +%H%M%S) | ||
| 101 | +mkdir -p logs | ||
| 102 | +mkdir -p images | ||
| 103 | +torchrun $DISTRIBUTED_ARGS posttrain_flux_dancegrpo.py \ | ||
| 104 | + $GPT_ARGS \ | ||
| 105 | + $MM_ARGS \ | ||
| 106 | + $OUTPUT_ARGS \ | ||
| 107 | + $GRPO_ARGS \ | ||
| 108 | + --distributed-backend nccl \ | ||
| 109 | + 2>&1 | tee logs/train_${logfile}.log | ||
| 110 | +chmod 440 logs/train_${logfile}.log | ||
| 111 | +find $SAVE_PATH -type d -exec chmod 750 {} \; | ||
| 112 | +find $SAVE_PATH -type f -exec chmod 640 {} \; | ||
| 113 | +STEP_TIME=`grep "step_time=" logs/train_${logfile}.log | awk -F '=' '{print$3}' | awk -F 's,' '{print$1}' | head -n 300 | tail -n 150 | awk '{sum+=$1} END {if (NR != 0) printf("%.1f",sum/NR)}'` | ||
| 114 | +GBS=`grep "Total train batch size" logs/train_${logfile}.log | awk -F '=' '{print$2}'` | ||
| 115 | +SAMPLES_PER_SECOND=$(awk -v gbs="${GBS}" -v step="${STEP_TIME}" 'BEGIN{printf "%.3f\n", gbs/step}') | ||
| 116 | +echo "Elapsed Time Per iteration: $STEP_TIME" | ||
| 117 | +echo "Average Samples per Second: $SAMPLES_PER_SECOND" | ||
| @@ -10,6 +10,7 @@ | |||
| 10 | # This modified file is released under the same license. | 10 | # This modified file is released under the same license. |
| 11 | 11 | ||
| 12 | import os | 12 | import os |
| 13 | +from concurrent.futures import ThreadPoolExecutor | ||
| 13 | 14 | ||
| 14 | import torch | 15 | import torch |
| 15 | import torch.distributed as dist | 16 | import torch.distributed as dist |
| @@ -19,8 +20,19 @@ from tqdm.auto import tqdm | |||
| 19 | from mindspeed_mm.tasks.rl.soragrpo.sora_grpo_trainer import SoraGRPOTrainer | 20 | from mindspeed_mm.tasks.rl.soragrpo.sora_grpo_trainer import SoraGRPOTrainer |
| 20 | from mindspeed_mm.tasks.rl.soragrpo.flux_grpo_model import FluxGRPOModel | 21 | from mindspeed_mm.tasks.rl.soragrpo.flux_grpo_model import FluxGRPOModel |
| 21 | 22 | ||
| 22 | - | ||
| 23 | class FluxGRPOTrainer(SoraGRPOTrainer): | 23 | class FluxGRPOTrainer(SoraGRPOTrainer): |
| 24 | + def __init__(self, train_valid_test_dataset_provider): | ||
| 25 | + super().__init__(train_valid_test_dataset_provider) | ||
| 26 | + self.step_counter = 0 | ||
| 27 | + self._image_save_executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="image_saver") | ||
| 28 | + | ||
| 29 | + def train(self): | ||
| 30 | + try: | ||
| 31 | + super().train() | ||
| 32 | + finally: | ||
| 33 | + self._image_save_executor.shutdown(wait=True) | ||
| 34 | + print("Image save executor has been shut down") | ||
| 35 | + | ||
| 24 | def model_provider(self, args): | 36 | def model_provider(self, args): |
| 25 | return FluxGRPOModel(args, device=self.device) | 37 | return FluxGRPOModel(args, device=self.device) |
| 26 | 38 | ||
| @@ -157,6 +169,7 @@ class FluxGRPOTrainer(SoraGRPOTrainer): | |||
| 157 | # dict of lists -> list of dicts for easier iteration | 169 | # dict of lists -> list of dicts for easier iteration |
| 158 | samples_batched_list = [dict(zip(samples_batched, x)) for x in zip(*samples_batched.values())] | 170 | samples_batched_list = [dict(zip(samples_batched, x)) for x in zip(*samples_batched.values())] |
| 159 | train_timesteps = int(len(samples["timesteps"][0]) * args.timestep_fraction) | 171 | train_timesteps = int(len(samples["timesteps"][0]) * args.timestep_fraction) |
| 172 | + self.step_counter += 1 | ||
| 160 | return samples_batched_list, train_timesteps, sigma_schedule, perms | 173 | return samples_batched_list, train_timesteps, sigma_schedule, perms |
| 161 | 174 | ||
| 162 | def sample_reference_model(self, args, caption, encoder_hidden_states, pooled_prompt_embeds, text_ids): | 175 | def sample_reference_model(self, args, caption, encoder_hidden_states, pooled_prompt_embeds, text_ids): |
| @@ -236,8 +249,10 @@ class FluxGRPOTrainer(SoraGRPOTrainer): | |||
| 236 | image_processor = VaeImageProcessor(16) | 249 | image_processor = VaeImageProcessor(16) |
| 237 | batch_decoded_images = image_processor.postprocess(image) | 250 | batch_decoded_images = image_processor.postprocess(image) |
| 238 | 251 | ||
| 239 | - for idx, image in zip(batch_idx, batch_decoded_images): | 252 | + if args.save_images: |
| 240 | - image.save(f"./images/flux_{rank}_{idx}.png") | 253 | + for idx, image in zip(batch_idx, batch_decoded_images): |
| 254 | + os.makedirs(f"./images/flux_grpo/step{self.step_counter}", exist_ok=True) | ||
| 255 | + self._image_save_executor.submit(self._image_save_async, image, f"./images/flux_grpo/step{self.step_counter}/rank{rank}_idx{idx}.png") | ||
| 241 | 256 | ||
| 242 | batch_caption = [caption[i] for i in batch_idx] | 257 | batch_caption = [caption[i] for i in batch_idx] |
| 243 | if args.use_hpsv2: | 258 | if args.use_hpsv2: |
| @@ -344,3 +359,10 @@ class FluxGRPOTrainer(SoraGRPOTrainer): | |||
| 344 | ) | 359 | ) |
| 345 | 360 | ||
| 346 | return latent_image_ids.to(device=device, dtype=dtype) | 361 | return latent_image_ids.to(device=device, dtype=dtype) |
| 362 | + | ||
| 363 | + | ||
| 364 | + def _save_image_async(image, filepath): | ||
| 365 | + try: | ||
| 366 | + image.save(filepath) | ||
| 367 | + except Exception as e: | ||
| 368 | + print(f"Error saving image {filepath}: {e}") | ||
| @@ -32,6 +32,7 @@ from mindspeed_mm.tasks.rl.soragrpo.utils.communications_flux import sp_parallel | |||
| 32 | from mindspeed_mm.tasks.rl.soragrpo.utils.fsdp_util import get_dit_fsdp_kwargs, apply_fsdp_checkpointing | 32 | from mindspeed_mm.tasks.rl.soragrpo.utils.fsdp_util import get_dit_fsdp_kwargs, apply_fsdp_checkpointing |
| 33 | from mindspeed_mm.tasks.rl.soragrpo.utils.parallel_states import initialize_sequence_parallel_state, \ | 33 | from mindspeed_mm.tasks.rl.soragrpo.utils.parallel_states import initialize_sequence_parallel_state, \ |
| 34 | get_sequence_parallel_state, destroy_sequence_parallel_group | 34 | get_sequence_parallel_state, destroy_sequence_parallel_group |
| 35 | +from mindspeed_mm.tasks.rl.soragrpo.utils.device import check_npu_version, NPUVersion | ||
| 35 | 36 | ||
| 36 | 37 | ||
| 37 | class SoraGRPOTrainer(ABC): | 38 | class SoraGRPOTrainer(ABC): |
| @@ -69,7 +70,9 @@ class SoraGRPOTrainer(ABC): | |||
| 69 | if world_size <= 8: | 70 | if world_size <= 8: |
| 70 | os.environ['HCCL_DETERMINISTIC'] = 'true' | 71 | os.environ['HCCL_DETERMINISTIC'] = 'true' |
| 71 | elif world_size >= 16: | 72 | elif world_size >= 16: |
| 72 | - os.environ['HCCL_OP_EXPANSION_MODE'] = 'AIV' | 73 | + # 950服务器不需要设置 HCCL_OP_EXPANSION_MODE=AIV |
| 74 | + if not check_npu_version(min_version=NPUVersion.A5): | ||
| 75 | + os.environ['HCCL_OP_EXPANSION_MODE'] = 'AIV' | ||
| 73 | 76 | ||
| 74 | # We use different seeds for the noise generation in each process to ensure that the noise is different in a batch. | 77 | # We use different seeds for the noise generation in each process to ensure that the noise is different in a batch. |
| 75 | if args.seed is not None: | 78 | if args.seed is not None: |
| @@ -240,6 +243,9 @@ class SoraGRPOTrainer(ABC): | |||
| 240 | print(f"rank {torch.distributed.get_rank()} profile complete") | 243 | print(f"rank {torch.distributed.get_rank()} profile complete") |
| 241 | 244 | ||
| 242 | step_time = time.time() - start_time | 245 | step_time = time.time() - start_time |
| 246 | + if torch.distributed.get_rank() == 0: | ||
| 247 | + print('===========================') | ||
| 248 | + print('step_time:', step_time) | ||
| 243 | step_times.append(step_time) | 249 | step_times.append(step_time) |
| 244 | 250 | ||
| 245 | progress_bar.set_postfix( | 251 | progress_bar.set_postfix( |
| @@ -265,14 +271,25 @@ class SoraGRPOTrainer(ABC): | |||
| 265 | lr_scheduler = self.lr_scheduler | 271 | lr_scheduler = self.lr_scheduler |
| 266 | 272 | ||
| 267 | total_loss = 0.0 | 273 | total_loss = 0.0 |
| 274 | + total_loss_addnum = 0 | ||
| 268 | optimizer.zero_grad() | 275 | optimizer.zero_grad() |
| 269 | 276 | ||
| 277 | + # Measure sample_reference (rollout) time | ||
| 278 | + rollout_start_time = time.time() | ||
| 270 | samples_batched_list, train_timesteps, sigma_schedule, perms = self.sample_reference(dataloader) | 279 | samples_batched_list, train_timesteps, sigma_schedule, perms = self.sample_reference(dataloader) |
| 280 | + rollout_time = time.time() - rollout_start_time | ||
| 281 | + | ||
| 282 | + # Initialize sft time tracking | ||
| 283 | + total_sft_time = 0.0 | ||
| 284 | + num_sft_calls = 0 | ||
| 271 | 285 | ||
| 272 | for i, sample in list(enumerate(samples_batched_list)): | 286 | for i, sample in list(enumerate(samples_batched_list)): |
| 273 | for j in range(train_timesteps): | 287 | for j in range(train_timesteps): |
| 274 | clip_range = args.clip_range | 288 | clip_range = args.clip_range |
| 275 | adv_clip_max = args.adv_clip_max | 289 | adv_clip_max = args.adv_clip_max |
| 290 | + | ||
| 291 | + # Measure sft time including grpo_one_step, loss calculation and backward | ||
| 292 | + sft_start_time = time.time() | ||
| 276 | new_log_probs = self.grpo_one_step( | 293 | new_log_probs = self.grpo_one_step( |
| 277 | sample, | 294 | sample, |
| 278 | perms[i][j], | 295 | perms[i][j], |
| @@ -296,9 +313,13 @@ class SoraGRPOTrainer(ABC): | |||
| 296 | args.gradient_accumulation_steps * train_timesteps) | 313 | args.gradient_accumulation_steps * train_timesteps) |
| 297 | 314 | ||
| 298 | loss.backward() | 315 | loss.backward() |
| 316 | + sft_time = time.time() - sft_start_time | ||
| 317 | + total_sft_time += sft_time | ||
| 318 | + num_sft_calls += 1 | ||
| 299 | avg_loss = loss.detach().clone() | 319 | avg_loss = loss.detach().clone() |
| 300 | dist.all_reduce(avg_loss, op=dist.ReduceOp.AVG) | 320 | dist.all_reduce(avg_loss, op=dist.ReduceOp.AVG) |
| 301 | total_loss += avg_loss.item() | 321 | total_loss += avg_loss.item() |
| 322 | + total_loss_addnum += 1 | ||
| 302 | 323 | ||
| 303 | if dist.get_rank() % self.world_size == 0: | 324 | if dist.get_rank() % self.world_size == 0: |
| 304 | print("hps reward", sample["rewards"].item()) | 325 | print("hps reward", sample["rewards"].item()) |
| @@ -312,6 +333,18 @@ class SoraGRPOTrainer(ABC): | |||
| 312 | lr_scheduler.step() | 333 | lr_scheduler.step() |
| 313 | optimizer.zero_grad() | 334 | optimizer.zero_grad() |
| 314 | dist.barrier() | 335 | dist.barrier() |
| 336 | + | ||
| 337 | + # Calculate average sft time | ||
| 338 | + avg_sft_time = total_sft_time / num_sft_calls if num_sft_calls > 0 else 0 | ||
| 339 | + if dist.get_rank() == 0: | ||
| 340 | + print("train step total_loss(sum)=", total_loss) | ||
| 341 | + print("----------------loss mean---------------") | ||
| 342 | + print("train step total_loss=", total_loss / total_loss_addnum) | ||
| 343 | + print("----------------time statistics---------------") | ||
| 344 | + print(f"rollout time: {rollout_time:.4f} seconds") | ||
| 345 | + print(f"average sft time: {avg_sft_time:.4f} seconds") | ||
| 346 | + print(f"total sft time: {total_sft_time:.4f} seconds") | ||
| 347 | + print(f"number of sft calls: {num_sft_calls}") | ||
| 315 | return total_loss, grad_norm.item() | 348 | return total_loss, grad_norm.item() |
| 316 | 349 | ||
| 317 | 350 | ||
| @@ -613,6 +646,12 @@ class SoraGRPOTrainer(ABC): | |||
| 613 | default=8, | 646 | default=8, |
| 614 | help="load rank batch size", | 647 | help="load rank batch size", |
| 615 | ) | 648 | ) |
| 649 | + parser.add_argument( | ||
| 650 | + "--save_images", | ||
| 651 | + action=argparse.BooleanOptionalAction, | ||
| 652 | + default=True, | ||
| 653 | + help="whether to save sampled images during training; disable for pure perf testing", | ||
| 654 | + ) | ||
| 616 | parser = mm_extra_args_provider(parser) | 655 | parser = mm_extra_args_provider(parser) |
| 617 | return parser.parse_args() | 656 | return parser.parse_args() |
| 618 | 657 | ||
| @@ -0,0 +1,43 @@ | |||
| 1 | +# Copyright (c) 2024, Huawei Technologies Co., Ltd. All rights reserved. | ||
| 2 | +from enum import IntEnum, unique | ||
| 3 | + | ||
| 4 | +from functools import lru_cache | ||
| 5 | +import torch_npu | ||
| 6 | + | ||
| 7 | + | ||
| 8 | + | ||
| 9 | +class NPUVersion(IntEnum): | ||
| 10 | + NONE = 0 | ||
| 11 | + A2 = 2 | ||
| 12 | + A3 = 3 | ||
| 13 | + A5 = 5 | ||
| 14 | + MAX_VERSION = 999 | ||
| 15 | + | ||
| 16 | + | ||
| 17 | +def check_npu_version(min_version=None, max_version=None): | ||
| 18 | + version = get_npu_version() | ||
| 19 | + if version == NPUVersion.NONE: | ||
| 20 | + return False | ||
| 21 | + if min_version is not None: | ||
| 22 | + if version < min_version: | ||
| 23 | + return False | ||
| 24 | + if max_version is not None: | ||
| 25 | + if version > max_version: | ||
| 26 | + return False | ||
| 27 | + return True | ||
| 28 | + | ||
| 29 | + | ||
| 30 | + | ||
| 31 | +def get_npu_version(): | ||
| 32 | + try: | ||
| 33 | + device_name = torch_npu.npu.get_device_name() | ||
| 34 | + except Exception: | ||
| 35 | + return NPUVersion.NONE | ||
| 36 | + if "Ascend910_95" in device_name or "Ascend950" in device_name: | ||
| 37 | + return NPUVersion.A5 | ||
| 38 | + elif "Ascend910_93" in device_name: | ||
| 39 | + return NPUVersion.A3 | ||
| 40 | + elif "Ascend910B" in device_name or "A2G" in device_name: | ||
| 41 | + return NPUVersion.A2 | ||
| 42 | + else: | ||
| 43 | + return NPUVersion.MAX_VERSION | ||
| @@ -114,6 +114,7 @@ class TestFluxGRPOTrainer(unittest.TestCase): | |||
| 114 | trainer.args.shift = 3.0 | 114 | trainer.args.shift = 3.0 |
| 115 | trainer.args.use_hpsv2 = True | 115 | trainer.args.use_hpsv2 = True |
| 116 | trainer.args.num_generations = 12 | 116 | trainer.args.num_generations = 12 |
| 117 | + trainer.args.save_images = False | ||
| 117 | 118 | ||
| 118 | trainer.hyper_model = MagicMock() | 119 | trainer.hyper_model = MagicMock() |
| 119 | trainer.hyper_model.diffuser = MagicMock() | 120 | trainer.hyper_model.diffuser = MagicMock() |
这个路径是否可以支持配置