Coverage for wrapper/flux/run_flux_benchmark.py: 19%

1# Flux inference with USP 

2 

3 

4import time 

5import torch 

6 

7from diffusers import FluxPipeline 

8 

9from flux_xfuser import parallelize_transformer 

10 

11from xfuser import xFuserArgs 

12from xfuser.config import FlexibleArgumentParser 

13from xfuser.core.distributed import ( 

14 get_world_group, 

15 get_data_parallel_world_size, 

16 get_data_parallel_rank, 

17 get_runtime_state, 

18 is_dp_last_group, 

19 initialize_runtime_state, 

20 get_pipeline_parallel_world_size, 

21) 

22 

23 

24def main() -> None: 

25 parser = FlexibleArgumentParser(description="xFuser Arguments") 

26 args = xFuserArgs.add_cli_args(parser).parse_args() 

27 engine_args = xFuserArgs.from_cli_args(args) 

28 engine_config, input_config = engine_args.create_config() 

29 engine_config.runtime_config.dtype = torch.bfloat16 

30 local_rank = get_world_group().local_rank 

31 

32 GPU = torch.cuda.get_device_name(0) 

33 

34 assert engine_args.pipefusion_parallel_degree == 1, "This script does not support PipeFusion." 

35 

36 # [-h] [--model MODEL] [--download-dir DOWNLOAD_DIR] [--trust-remote-code] [--warmup_steps WARMUP_STEPS] 

37 # [--use_parallel_vae] 

38 # [--use_torch_compile] [--use_onediff] 

39 # [--use_teacache] [--use_fbcache] [--use_ray] [--ray_world_size RAY_WORLD_SIZE] 

40 # [--dit_parallel_size DIT_PARALLEL_SIZE] 

41 # [--use_cfg_parallel] 

42 # [--data_parallel_degree DATA_PARALLEL_DEGREE] [--ulysses_degree ULYSSES_DEGREE] [--ring_degree RING_DEGREE] 

43 # [--pipefusion_parallel_degree PIPEFUSION_PARALLEL_DEGREE 

44 # [--num_pipeline_patch NUM_PIPELINE_PATCH] [--attn_layer_num_for_pp [ATTN_LAYER_NUM_FOR_PP ...]] 

45 # [--tensor_parallel_degree TENSOR_PARALLEL_DEGREE] 

46 # [--vae_parallel_size VAE_PARALLEL_SIZE] [--split_scheme SPLIT_SCHEME] [--height HEIGHT] [--width WIDTH] 

47 # [--num_frames NUM_FRAMES] [--img_file_path IMG_FILE_PATH] 

48 # [--prompt [PROMPT ...]] [--no_use_resolution_binning] [--negative_prompt [NEGATIVE_PROMPT ...]] 

49 # [--num_inference_steps NUM_INFERENCE_STEPS] 

50 # [--max_sequence_length MAX_SEQUENCE_LENGTH] [--seed SEED] [--output_type OUTPUT_TYPE] 

51 # [--enable_sequential_cpu_offload] [--enable_model_cpu_offload] [--enable_tiling] 

52 # [--enable_slicing] [--use_fp8_t5_encoder] [--use_fast_attn] [--n_calib N_CALIB] [--threshold THRESHOLD] 

53 # [--window_size WINDOW_SIZE] [--coco_path COCO_PATH] 

54 # [--use_cache] 

55 

56 cache_args = { 

57 "use_teacache": engine_args.use_teacache, 

58 "use_fbcache": engine_args.use_fbcache, 

59 "rel_l1_thresh": 0.12, 

60 "return_hidden_states_first": False, 

61 "num_steps": input_config.num_inference_steps, 

62 } 

63 

64 # pipe = xFuserFluxPipeline.from_pretrained( 

65 pipe = FluxPipeline.from_pretrained( 

66 pretrained_model_name_or_path=engine_config.model_config.model, 

67 engine_config=engine_config, 

68 cache_args=cache_args, 

69 torch_dtype=torch.bfloat16, 

70 ) 

71 pipe = pipe.to(f"cuda:{local_rank}") 

72 

73 parameter_peak_memory = torch.cuda.max_memory_allocated(device=f"cuda:{local_rank}") 

74 

75 initialize_runtime_state(pipe, engine_config) 

76 get_runtime_state().set_input_parameters( 

77 height=input_config.height, 

78 width=input_config.width, 

79 batch_size=1, 

80 num_inference_steps=input_config.num_inference_steps, 

81 max_condition_sequence_length=512, 

82 split_text_embed_in_sp=get_pipeline_parallel_world_size() == 1, 

83 ) 

84 

85 parallelize_transformer(pipe) 

86 

87 if engine_config.runtime_config.use_torch_compile: 

88 torch._inductor.config.reorder_for_compute_comm_overlap = True 

89 pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune-no-cudagraphs") 

90 

91 # one step to warmup the torch compiler 

92 output = pipe( 

93 height=input_config.height, 

94 width=input_config.width, 

95 prompt=input_config.prompt, 

96 num_inference_steps=1, 

97 output_type=input_config.output_type, 

98 generator=torch.Generator(device="cuda").manual_seed(input_config.seed), 

99 ).images 

100 

101 # Run to warm up the model 

102 output = pipe( 

103 height=input_config.height, 

104 width=input_config.width, 

105 prompt="warmup prompt", 

106 num_inference_steps=1, 

107 output_type=input_config.output_type, 

108 generator=torch.Generator(device="cuda").manual_seed(input_config.seed), 

109 ) 

110 

111 # Actual run 

112 torch.cuda.reset_peak_memory_stats() 

113 start_time = time.time() 

114 

115 input_config.prompt = "a photo of an astronaut riding a horse on mars" 

116 output = pipe( 

117 height=input_config.height, 

118 width=input_config.width, 

119 prompt=input_config.prompt, 

120 num_inference_steps=input_config.num_inference_steps, 

121 output_type=input_config.output_type, 

122 generator=torch.Generator(device="cuda").manual_seed(input_config.seed), 

123 ) 

124 torch.cuda.synchronize() # Ensure all above CUDA ops are done 

125 end_time = time.time() 

126 elapsed_time = end_time - start_time 

127 peak_memory = torch.cuda.max_memory_allocated(device=f"cuda:{local_rank}") 

128 

129 parallel_info = ( 

130 f"dp{engine_args.data_parallel_degree}_cfg{engine_config.parallel_config.cfg_degree}_" 

131 f"ulysses{engine_args.ulysses_degree}_ring{engine_args.ring_degree}_" 

132 f"tp{engine_args.tensor_parallel_degree}_" 

133 f"pp{engine_args.pipefusion_parallel_degree}_patch{engine_args.num_pipeline_patch}" 

134 ) 

135 if input_config.output_type == "pil": 

136 dp_group_index = get_data_parallel_rank() 

137 num_dp_groups = get_data_parallel_world_size() 

138 dp_batch_size = (input_config.batch_size + num_dp_groups - 1) // num_dp_groups 

139 if is_dp_last_group(): 

140 for i, image in enumerate(output.images): 

141 image_rank = dp_group_index * dp_batch_size + i 

142 image_name = f"flux_result_{parallel_info}_{image_rank}_tc_{engine_args.use_torch_compile}.png" 

143 image.save(f"./results/{image_name}") 

144 print(f"image {i} saved to ./results/{image_name}") 

145 

146 # Write into a file 

147 if get_world_group().rank == 0: 

148 with open("flux_parallel_result.csv", "a") as f: 

149 dp = engine_args.data_parallel_degree 

150 up = engine_args.ulysses_degree 

151 rp = engine_args.ring_degree 

152 tp = engine_args.tensor_parallel_degree 

153 cfg = engine_config.parallel_config.cfg_degree 

154 torchcompile = engine_config.runtime_config.use_torch_compile 

155 teacache = engine_args.use_teacache 

156 fbcache = engine_args.use_fbcache 

157 steps = input_config.num_inference_steps 

158 height = input_config.height 

159 width = input_config.width 

160 f.write(f"{GPU},{get_world_group().world_size},{dp},{up},{rp},{tp},{cfg}," 

161 f"{torchcompile},{teacache},{fbcache},{steps},{height},{width},{elapsed_time:.2f}\n") 

162 

163 if get_world_group().rank == 0: 

164 print( 

165 f"epoch time: {elapsed_time:.2f} sec, parameter memory: {parameter_peak_memory / 1e9:.2f} GB, " 

166 f"memory: {peak_memory / 1e9:.2f} GB" 

167 ) 

168 print(f"VAE: {pipe.vae.elapsed_time:.2f} sec") 

169 print(f"Transformer: {pipe.transformer.get_elapsed_time():.2f} sec") 

170 print(f"Scheduler: {pipe.scheduler.elapsed_time:.2f} sec") 

171 # get_runtime_state().destroy_distributed_env() 

172 get_runtime_state().destory_distributed_env() 

173 

174 

175if __name__ == "__main__": 

176 main()