[run.sh] distributed_sharded auto-set NUM_PROCESSES=1 (all visible GPUs)
[run.sh] Launch mode=distributed_sharded (DeepSpeed ZeRO-3)
2026-04-07 13:37:34.662042: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2026-04-07 13:37:37.839913: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 AVX512_FP16 AVX_VNNI AMX_TILE AMX_INT8 AMX_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2026-04-07 13:37:42.467616: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2026-04-07 13:37:42.473892: I external/local_xla/xla/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
[train.py] Disabled flash/mem-efficient SDP kernels; using math SDP backend.
[2026-04-07 13:38:10,838][root][INFO] - /g/data/rr81/aev/bin/x86_64-conda-linux-gnu-cc -march=nocona -mtune=haswell -ftree-vectorize -fPIC -fstack-protector-strong -fno-plt -O2 -ffunction-sections -pipe -isystem /g/data/rr81/aev/include -I/g/data/rr81/aev/targets/x86_64-linux/include -I/g/data/rr81/aev/targets/x86_64-linux/include/cccl -DNDEBUG -D_FORTIFY_SOURCE=2 -O2 -isystem /g/data/rr81/aev/include -I/g/data/rr81/aev/targets/x86_64-linux/include -I/g/data/rr81/aev/targets/x86_64-linux/include/cccl -fPIC -march=nocona -mtune=haswell -ftree-vectorize -fPIC -fstack-protector-strong -fno-plt -O2 -ffunction-sections -pipe -isystem /g/data/rr81/aev/include -I/g/data/rr81/aev/targets/x86_64-linux/include -I/g/data/rr81/aev/targets/x86_64-linux/include/cccl -c /scratch/rr81/ma5430/tmp/tmpspb_o7mh/test.c -o /scratch/rr81/ma5430/tmp/tmpspb_o7mh/test.o
[2026-04-07 13:38:24,727][root][INFO] - /g/data/rr81/aev/bin/x86_64-conda-linux-gnu-cc -Wl,-O2 -Wl,--sort-common -Wl,--as-needed -Wl,-z,relro -Wl,-z,now -Wl,--disable-new-dtags -Wl,--gc-sections -Wl,-rpath,/g/data/rr81/aev/lib -Wl,-rpath-link,/g/data/rr81/aev/lib -L/g/data/rr81/aev/lib -L/g/data/rr81/aev/targets/x86_64-linux/lib -L/g/data/rr81/aev/targets/x86_64-linux/lib/stubs /scratch/rr81/ma5430/tmp/tmpspb_o7mh/test.o -laio -o /scratch/rr81/ma5430/tmp/tmpspb_o7mh/a.out
[2026-04-07 13:38:25,328][root][INFO] - /g/data/rr81/aev/bin/x86_64-conda-linux-gnu-cc -march=nocona -mtune=haswell -ftree-vectorize -fPIC -fstack-protector-strong -fno-plt -O2 -ffunction-sections -pipe -isystem /g/data/rr81/aev/include -I/g/data/rr81/aev/targets/x86_64-linux/include -I/g/data/rr81/aev/targets/x86_64-linux/include/cccl -DNDEBUG -D_FORTIFY_SOURCE=2 -O2 -isystem /g/data/rr81/aev/include -I/g/data/rr81/aev/targets/x86_64-linux/include -I/g/data/rr81/aev/targets/x86_64-linux/include/cccl -fPIC -march=nocona -mtune=haswell -ftree-vectorize -fPIC -fstack-protector-strong -fno-plt -O2 -ffunction-sections -pipe -isystem /g/data/rr81/aev/include -I/g/data/rr81/aev/targets/x86_64-linux/include -I/g/data/rr81/aev/targets/x86_64-linux/include/cccl -c /scratch/rr81/ma5430/tmp/tmp6efs5r08/test.c -o /scratch/rr81/ma5430/tmp/tmp6efs5r08/test.o
[2026-04-07 13:38:25,385][root][INFO] - /g/data/rr81/aev/bin/x86_64-conda-linux-gnu-cc -Wl,-O2 -Wl,--sort-common -Wl,--as-needed -Wl,-z,relro -Wl,-z,now -Wl,--disable-new-dtags -Wl,--gc-sections -Wl,-rpath,/g/data/rr81/aev/lib -Wl,-rpath-link,/g/data/rr81/aev/lib -L/g/data/rr81/aev/lib -L/g/data/rr81/aev/targets/x86_64-linux/lib -L/g/data/rr81/aev/targets/x86_64-linux/lib/stubs /scratch/rr81/ma5430/tmp/tmp6efs5r08/test.o -L/g/data/rr81/aev -L/g/data/rr81/aev/lib64 -lcufile -o /scratch/rr81/ma5430/tmp/tmp6efs5r08/a.out
[2026-04-07 13:38:25,551][root][INFO] - /g/data/rr81/aev/bin/x86_64-conda-linux-gnu-cc -march=nocona -mtune=haswell -ftree-vectorize -fPIC -fstack-protector-strong -fno-plt -O2 -ffunction-sections -pipe -isystem /g/data/rr81/aev/include -I/g/data/rr81/aev/targets/x86_64-linux/include -I/g/data/rr81/aev/targets/x86_64-linux/include/cccl -DNDEBUG -D_FORTIFY_SOURCE=2 -O2 -isystem /g/data/rr81/aev/include -I/g/data/rr81/aev/targets/x86_64-linux/include -I/g/data/rr81/aev/targets/x86_64-linux/include/cccl -fPIC -march=nocona -mtune=haswell -ftree-vectorize -fPIC -fstack-protector-strong -fno-plt -O2 -ffunction-sections -pipe -isystem /g/data/rr81/aev/include -I/g/data/rr81/aev/targets/x86_64-linux/include -I/g/data/rr81/aev/targets/x86_64-linux/include/cccl -c /scratch/rr81/ma5430/tmp/tmp0ddzq53f/test.c -o /scratch/rr81/ma5430/tmp/tmp0ddzq53f/test.o
[2026-04-07 13:38:25,609][root][INFO] - /g/data/rr81/aev/bin/x86_64-conda-linux-gnu-cc -Wl,-O2 -Wl,--sort-common -Wl,--as-needed -Wl,-z,relro -Wl,-z,now -Wl,--disable-new-dtags -Wl,--gc-sections -Wl,-rpath,/g/data/rr81/aev/lib -Wl,-rpath-link,/g/data/rr81/aev/lib -L/g/data/rr81/aev/lib -L/g/data/rr81/aev/targets/x86_64-linux/lib -L/g/data/rr81/aev/targets/x86_64-linux/lib/stubs /scratch/rr81/ma5430/tmp/tmp0ddzq53f/test.o -laio -o /scratch/rr81/ma5430/tmp/tmp0ddzq53f/a.out
[2026-04-07 13:38:28,753][accelerate.utils.other][WARNING] - Detected kernel version 4.18.0, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.
[2026-04-07 13:38:28,754][trainer.accelerators.base_accelerator][INFO] - Setting seed 42
[2026-04-07 13:38:28,771][trainer.accelerators.base_accelerator][INFO] - Initialized accelerator: rank=0
[2026-04-07 13:38:28,777][__main__][INFO] - Config can be found in logs/v5/reward_model/step_sana_sana_sprint_0_6b_1024_variable-t_lr1e-5_step-8000_filter2_time951/config.yaml
[2026-04-07 13:38:28,778][__main__][INFO] - Loading task
[2026-04-07 13:38:29,730][__main__][INFO] - Loading model
`torch_dtype` is deprecated! Use `dtype` instead!

Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]
Loading checkpoint shards:  50%|█████     | 1/2 [00:01<00:01,  1.82s/it]
Loading checkpoint shards: 100%|██████████| 2/2 [00:01<00:00,  1.22it/s]
Loading checkpoint shards: 100%|██████████| 2/2 [00:01<00:00,  1.03it/s]
[2026-04-07 13:38:33,300][__main__][INFO] - Loading criterion
[2026-04-07 13:38:33,301][__main__][INFO] - Loading optimizer
[2026-04-07 13:38:33,305][__main__][INFO] - Loading lr scheduler
[2026-04-07 13:38:33,306][__main__][INFO] - Loading dataloaders
[2026-04-07 13:38:33,307][trainer.datasets.step_sana_hf_dataset][INFO] - Using step-aware datasets
[2026-04-07 13:38:33,307][trainer.datasets.step_sana_hf_dataset][INFO] - Loading train dataset
[2026-04-07 13:38:33,307][trainer.datasets.step_sana_hf_dataset][INFO] - Batch size is 4
[2026-04-07 13:38:33,309][trainer.datasets.step_sana_hf_dataset][INFO] - Loading cached offline split 'train' from 387 parquet shards
[2026-04-07 13:38:34,488][trainer.datasets.step_sana_hf_dataset][INFO] - Loaded 583747 examples from train dataset
[2026-04-07 13:38:35,043][trainer.datasets.step_sana_hf_dataset][INFO] - Keeping only examples with pesudo preference, filter_strategy: 2
[2026-04-07 13:38:35,050][trainer.datasets.step_sana_hf_dataset][INFO] - Loaded 583747 examples from train dataset
[2026-04-07 13:38:35,241][trainer.datasets.step_sana_hf_dataset][INFO] - Kept 177076 examples from train dataset
[2026-04-07 13:38:35,242][trainer.datasets.step_sana_hf_dataset][INFO] - Loaded 177076 examples from train dataset
[2026-04-07 13:38:36,181][trainer.datasets.step_sana_hf_dataset][INFO] - Using step-aware datasets
[2026-04-07 13:38:36,181][trainer.datasets.step_sana_hf_dataset][INFO] - Loading validation_unique dataset
[2026-04-07 13:38:36,181][trainer.datasets.step_sana_hf_dataset][INFO] - Batch size is 4
[2026-04-07 13:38:36,182][trainer.datasets.step_sana_hf_dataset][INFO] - Loading cached offline split 'validation_unique' from 1 parquet shards
[2026-04-07 13:38:36,193][trainer.datasets.step_sana_hf_dataset][INFO] - Loaded 500 examples from validation_unique dataset
[2026-04-07 13:38:36,194][trainer.datasets.step_sana_hf_dataset][INFO] - Keeping only examples with label in validation_unique split

Filter:   0%|          | 0/500 [00:00<?, ? examples/s]
Filter: 100%|██████████| 500/500 [00:00<00:00, 3601.50 examples/s]
Filter: 100%|██████████| 500/500 [00:00<00:00, 3530.07 examples/s]
[2026-04-07 13:38:36,359][trainer.datasets.step_sana_hf_dataset][INFO] - Kept 425 examples from validation_unique dataset
[2026-04-07 13:38:36,360][trainer.datasets.step_sana_hf_dataset][INFO] - Loaded 425 examples from validation_unique dataset
[2026-04-07 13:38:37,357][trainer.datasets.step_sana_hf_dataset][INFO] - Using step-aware datasets
[2026-04-07 13:38:37,357][trainer.datasets.step_sana_hf_dataset][INFO] - Loading test_unique dataset
[2026-04-07 13:38:37,357][trainer.datasets.step_sana_hf_dataset][INFO] - Batch size is 4
[2026-04-07 13:38:37,358][trainer.datasets.step_sana_hf_dataset][INFO] - Loading cached offline split 'test_unique' from 1 parquet shards
[2026-04-07 13:38:37,369][trainer.datasets.step_sana_hf_dataset][INFO] - Loaded 500 examples from test_unique dataset
[2026-04-07 13:38:37,369][trainer.datasets.step_sana_hf_dataset][INFO] - Keeping only examples with label in test_unique split

Filter:   0%|          | 0/500 [00:00<?, ? examples/s]
Filter: 100%|██████████| 500/500 [00:00<00:00, 3381.00 examples/s]
Filter: 100%|██████████| 500/500 [00:00<00:00, 3311.47 examples/s]
[2026-04-07 13:38:37,545][trainer.datasets.step_sana_hf_dataset][INFO] - Kept 432 examples from test_unique dataset
[2026-04-07 13:38:37,545][trainer.datasets.step_sana_hf_dataset][INFO] - Loaded 432 examples from test_unique dataset
[2026-04-07 13:38:38,482][accelerate.accelerator][INFO] - Since you passed both train and evaluation dataloader, `is_train_batch_min` (here True will decide the `train_batch_size` (4).
[2026-04-07 13:38:44,136] [WARNING] [lr_schedules.py:690:get_lr] Attempting to get learning rate from scheduler before it has started
[2026-04-07 13:38:44,144][trainer.accelerators.base_accelerator][INFO] - num_update_steps_per_epoch = 44269
[2026-04-07 13:38:44,144][trainer.accelerators.base_accelerator][INFO] - num_batches = 44269
[2026-04-07 13:38:44,144][trainer.accelerators.base_accelerator][INFO] - num_epochs = 1
[2026-04-07 13:38:44,182][trainer.accelerators.base_accelerator][INFO] - Initializing trackers
[2026-04-07 13:38:44,182][trainer.accelerators.base_accelerator][INFO] - Training config:
CONFIG
├── accelerator
│   └── _target_: trainer.accelerators.deepspeed_accelerator.DeepSpeedAccelerator                                                                                                
│       output_dir: logs/v5/reward_model/step_sana_sana_sprint_0_6b_1024_variable-t_lr1e-5_step-8000_filter2_time951                                                             
│       mixed_precision: BF16                                                                                                                                                    
│       gradient_accumulation_steps: 1                                                                                                                                           
│       log_with: null                                                                                                                                                           
│       debug:                                                                                                                                                                   
│         activate: false                                                                                                                                                        
│         port: 5900                                                                                                                                                             
│       seed: 42                                                                                                                                                                 
│       resume_from_checkpoint: false                                                                                                                                            
│       max_steps: 8000                                                                                                                                                          
│       num_epochs: 1                                                                                                                                                            
│       validate_steps: 100                                                                                                                                                      
│       generalization_validate_steps: 500                                                                                                                                       
│       eval_on_start: false                                                                                                                                                     
│       project_name: reward_model                                                                                                                                               
│       run_name: step_sana_sana_sprint_0_6b_1024_variable-t_lr1e-5_step-8000_filter2_time951                                                                                    
│       max_grad_norm: 1.0                                                                                                                                                       
│       save_steps: 100                                                                                                                                                          
│       metric_name: accuracy                                                                                                                                                    
│       metric_mode: MAX                                                                                                                                                         
│       limit_num_checkpoints: 1                                                                                                                                                 
│       save_only_if_best: true                                                                                                                                                  
│       dynamo_backend: 'NO'                                                                                                                                                     
│       keep_best_ckpts: true                                                                                                                                                    
│       progress_log_interval: 50                                                                                                                                                
│       deepspeed:                                                                                                                                                               
│         fp16:                                                                                                                                                                  
│           enabled: false                                                                                                                                                       
│         bf16:                                                                                                                                                                  
│           enabled: true                                                                                                                                                        
│         optimizer:                                                                                                                                                             
│           type: AdamW                                                                                                                                                          
│           params:                                                                                                                                                              
│             lr: auto                                                                                                                                                           
│             weight_decay: auto                                                                                                                                                 
│             torch_adam: true                                                                                                                                                   
│             adam_w_mode: true                                                                                                                                                  
│         scheduler:                                                                                                                                                             
│           type: WarmupDecayLR                                                                                                                                                  
│           params:                                                                                                                                                              
│             warmup_min_lr: auto                                                                                                                                                
│             warmup_max_lr: auto                                                                                                                                                
│             warmup_num_steps: auto                                                                                                                                             
│             total_num_steps: auto                                                                                                                                              
│         zero_optimization:                                                                                                                                                     
│           stage: 3                                                                                                                                                             
│           allgather_partitions: true                                                                                                                                           
│           allgather_bucket_size: 200000000.0                                                                                                                                   
│           overlap_comm: true                                                                                                                                                   
│           reduce_scatter: true                                                                                                                                                 
│           reduce_bucket_size: 500000000                                                                                                                                        
│           contiguous_gradients: true                                                                                                                                           
│         gradient_accumulation_steps: 1                                                                                                                                         
│         gradient_clipping: 1.0                                                                                                                                                 
│         steps_per_print: 1                                                                                                                                                     
│         train_batch_size: auto                                                                                                                                                 
│         train_micro_batch_size_per_gpu: auto                                                                                                                                   
│         wall_clock_breakdown: false                                                                                                                                            
│       deepspeed_final:                                                                                                                                                         
│         fp16:                                                                                                                                                                  
│           enabled: false                                                                                                                                                       
│         bf16:                                                                                                                                                                  
│           enabled: true                                                                                                                                                        
│         optimizer:                                                                                                                                                             
│           type: AdamW                                                                                                                                                          
│           params:                                                                                                                                                              
│             lr: auto                                                                                                                                                           
│             weight_decay: auto                                                                                                                                                 
│             torch_adam: true                                                                                                                                                   
│             adam_w_mode: true                                                                                                                                                  
│         scheduler:                                                                                                                                                             
│           type: WarmupDecayLR                                                                                                                                                  
│           params:                                                                                                                                                              
│             warmup_min_lr: auto                                                                                                                                                
│             warmup_max_lr: auto                                                                                                                                                
│             warmup_num_steps: auto                                                                                                                                             
│             total_num_steps: auto                                                                                                                                              
│         zero_optimization:                                                                                                                                                     
│           stage: 3                                                                                                                                                             
│           allgather_partitions: true                                                                                                                                           
│           allgather_bucket_size: 200000000.0                                                                                                                                   
│           overlap_comm: true                                                                                                                                                   
│           reduce_scatter: true                                                                                                                                                 
│           reduce_bucket_size: 500000000                                                                                                                                        
│           contiguous_gradients: true                                                                                                                                           
│         gradient_accumulation_steps: 1                                                                                                                                         
│         gradient_clipping: 1.0                                                                                                                                                 
│         steps_per_print: .inf                                                                                                                                                  
│         train_batch_size: auto                                                                                                                                                 
│         train_micro_batch_size_per_gpu: auto                                                                                                                                   
│         wall_clock_breakdown: false                                                                                                                                            
│                                                                                                                                                                                
├── task
│   └── limit_examples_to_wandb: 50                                                                                                                                              
│       _target_: trainer.tasks.step_sana_task.StepSanaTask                                                                                                                      
│       pretrained_model_name_or_path: Efficient-Large-Model/Sana_Sprint_0.6B_1024px_diffusers                                                                                   
│       tokenizer_subfolder: tokenizer                                                                                                                                           
│       label_0_column_name: label_0                                                                                                                                             
│       label_1_column_name: label_1                                                                                                                                             
│       input_ids_column_name: input_ids                                                                                                                                         
│       input_ids_2_column_name: input_ids_2                                                                                                                                     
│       pixels_0_column_name: pixel_values_0                                                                                                                                     
│       pixels_1_column_name: pixel_values_1                                                                                                                                     
│       timestep_column_name: timestep                                                                                                                                           
│       constant_timestep: 1                                                                                                                                                     
│                                                                                                                                                                                
├── model
│   └── _target_: trainer.models.sana_preference_model.SanaPreferenceModel                                                                                                       
│       pretrained_model_name_or_path: Efficient-Large-Model/Sana_Sprint_0.6B_1024px_diffusers                                                                                   
│       pretrained_vae_name_or_path: ''                                                                                                                                          
│       model_profile: sana_sprint_0_6b_1024                                                                                                                                     
│       projection_dim: 1024                                                                                                                                                     
│       logit_scale_init_value: 2.6592                                                                                                                                           
│       freeze_text_encoder: false                                                                                                                                               
│       guidance_scale: 7.5                                                                                                                                                      
│       noise_offset: false                                                                                                                                                      
│       noise_offset_coeff: 0.05                                                                                                                                                 
│       max_sequence_length: 300                                                                                                                                                 
│       max_sequence_length_2: 300                                                                                                                                               
│       image_size: 1024                                                                                                                                                         
│                                                                                                                                                                                
├── criterion
│   └── _target_: trainer.criterions.step_clip_criterion_sana.StepSanaCLIPCriterion                                                                                              
│       is_distributed: true                                                                                                                                                     
│       label_0_column_name: label_0                                                                                                                                             
│       label_1_column_name: label_1                                                                                                                                             
│       input_ids_column_name: input_ids                                                                                                                                         
│       input_ids_2_column_name: input_ids_2                                                                                                                                     
│       pixels_0_column_name: pixel_values_0                                                                                                                                     
│       pixels_1_column_name: pixel_values_1                                                                                                                                     
│       num_examples_per_prompt_column_name: num_example_per_prompt                                                                                                              
│       timestep_column_name: timestep                                                                                                                                           
│       loss_type: pair                                                                                                                                                          
│       batch_coeff: 1.0                                                                                                                                                         
│       aux_loss_coeff: 1.0                                                                                                                                                      
│                                                                                                                                                                                
├── dataset
│   └── train_split_name: train                                                                                                                                                  
│       valid_split_name: validation_unique                                                                                                                                      
│       test_split_name: test_unique                                                                                                                                             
│       batch_size: 4                                                                                                                                                            
│       num_workers: 2                                                                                                                                                           
│       drop_last: true                                                                                                                                                          
│       _target_: trainer.datasets.step_sana_hf_dataset.StepSanaHFDataset                                                                                                        
│       dataset_name: pickapic-anonymous/pickapic_v1                                                                                                                             
│       dataset_config_name: null                                                                                                                                                
│       from_disk: false                                                                                                                                                         
│       cache_dir: null                                                                                                                                                          
│       caption_column_name: caption                                                                                                                                             
│       input_ids_column_name: input_ids                                                                                                                                         
│       input_ids_2_column_name: input_ids_2                                                                                                                                     
│       image_0_column_name: jpg_0                                                                                                                                               
│       image_1_column_name: jpg_1                                                                                                                                               
│       label_0_column_name: label_0                                                                                                                                             
│       label_1_column_name: label_1                                                                                                                                             
│       are_different_column_name: are_different                                                                                                                                 
│       has_label_column_name: has_label                                                                                                                                         
│       pixels_0_column_name: pixel_values_0                                                                                                                                     
│       pixels_1_column_name: pixel_values_1                                                                                                                                     
│       timestep_column_name: timestep                                                                                                                                           
│       constant_timestep: 1                                                                                                                                                     
│       variable_timestep: true                                                                                                                                                  
│       largest_timestep: 951                                                                                                                                                    
│       compare_between_timestep: false                                                                                                                                          
│       timestep_comparison_column_name: timestep_comparison                                                                                                                     
│       timestep_interval: 1                                                                                                                                                     
│       num_examples_per_prompt_column_name: num_example_per_prompt                                                                                                              
│       keep_only_different: false                                                                                                                                               
│       keep_only_with_label: false                                                                                                                                              
│       keep_only_with_label_in_non_train: true                                                                                                                                  
│       keep_only_with_pesudo_preference: true                                                                                                                                   
│       pseudo_preference_path: /g/data/rr81/LPO/lrm/lrm_sana/vqa_aes_clip_score_mp.csv                                                                                          
│       filter_strategy: 2                                                                                                                                                       
│       processor:                                                                                                                                                               
│         pretrained_model_name_or_path: Efficient-Large-Model/Sana_Sprint_0.6B_1024px_diffusers                                                                                 
│         max_sequence_length: 300                                                                                                                                               
│         max_sequence_length_2: 300                                                                                                                                             
│         image_size: 1024                                                                                                                                                       
│         random_crop: false                                                                                                                                                     
│         no_hflip: true                                                                                                                                                         
│       limit_examples_per_prompt: -1                                                                                                                                            
│       only_on_best: false                                                                                                                                                      
│                                                                                                                                                                                
├── optimizer
│   └── _target_: trainer.optimizers.dummy_optimizer.BaseDummyOptim                                                                                                              
│       lr: 1.0e-05                                                                                                                                                              
│       weight_decay: 0.3                                                                                                                                                        
│                                                                                                                                                                                
├── lr_scheduler
│   └── _target_: trainer.lr_schedulers.dummy_lr_scheduler.instantiate_dummy_lr_scheduler                                                                                        
│       lr: 1.0e-05                                                                                                                                                              
│       lr_warmup_steps: 1000                                                                                                                                                    
│       total_num_steps: 8000                                                                                                                                                    
│                                                                                                                                                                                
├── debug
│   └── activate: false                                                                                                                                                          
│       port: 5900                                                                                                                                                               
│                                                                                                                                                                                
└── output_dir
    └── logs/v5/reward_model/step_sana_sana_sprint_0_6b_1024_variable-t_lr1e-5_step-8000_filter2_time951                                                                         
[2026-04-07 13:38:44,256][trainer.accelerators.base_accelerator][INFO] - nvidia-smi stats: {'gpu_0_mem_used_gb': 35.8818359375}
[2026-04-07 13:38:44,256][trainer.accelerators.base_accelerator][INFO] - ***** Running training *****
[2026-04-07 13:38:44,257][trainer.accelerators.base_accelerator][INFO] -   Instantaneous batch size per device = 4
[2026-04-07 13:38:44,257][trainer.accelerators.base_accelerator][INFO] -   Total train batch size (w. parallel, distributed & accumulation) = 4
[2026-04-07 13:38:44,257][trainer.accelerators.base_accelerator][INFO] -   Gradient Accumulation steps = 1
[2026-04-07 13:38:44,257][trainer.accelerators.base_accelerator][INFO] -   Total warmup steps = 1000
[2026-04-07 13:38:44,257][trainer.accelerators.base_accelerator][INFO] -   Total training steps = 8000
[2026-04-07 13:38:44,257][trainer.accelerators.base_accelerator][INFO] -   Total epochs = 1
[2026-04-07 13:38:44,258][trainer.accelerators.base_accelerator][INFO] -   Steps per epoch = 44269
[2026-04-07 13:38:44,258][trainer.accelerators.base_accelerator][INFO] -   Update steps per epoch = 44269
[2026-04-07 13:38:44,258][trainer.accelerators.base_accelerator][INFO] -   Total optimization steps = 8000
[2026-04-07 13:38:44,258][trainer.accelerators.base_accelerator][INFO] -   Mixed precision = bf16
[2026-04-07 13:38:44,258][trainer.accelerators.base_accelerator][INFO] -   World size = 1

  0%|          | 0/8000 [00:00<?, ?it/s]
Steps:   0%|          | 0/8000 [00:00<?, ?it/s][2026-04-07 13:38:44,260][__main__][INFO] - task: StepSanaTask
[2026-04-07 13:38:44,260][__main__][INFO] - model: DeepSpeedEngine
[2026-04-07 13:38:44,269][__main__][INFO] - num. model params: 0M
[2026-04-07 13:38:44,278][__main__][INFO] - num. model trainable params: 0M
[2026-04-07 13:38:44,278][__main__][INFO] - criterion: StepSanaCLIPCriterion
[2026-04-07 13:38:44,278][__main__][INFO] - num. train examples: 177076
[2026-04-07 13:38:44,279][__main__][INFO] - num. valid examples: 425
[2026-04-07 13:38:44,279][__main__][INFO] - num. test examples: 432
[2026-04-07 13:38:44,279][__main__][INFO] - ========== TRAIN LOOP START (eval_on_start=False, validate_steps=100, progress_log_interval=50) ==========
[2026-04-07 13:38:45,177] [WARNING] [torch_autocast.py:122:autocast_if_enabled] torch.autocast is enabled outside DeepSpeed but disabled within the DeepSpeed engine. If you are using DeepSpeed's built-in mixed precision, the engine will follow the settings in bf16/fp16 section. To use torch's native autocast instead, configure the `torch_autocast` section in the DeepSpeed config.
Error executing job with overrides: ['accelerator.mixed_precision=BF16', 'model.model_profile=sana_sprint_0_6b_1024', 'model.pretrained_model_name_or_path=Efficient-Large-Model/Sana_Sprint_0.6B_1024px_diffusers', 'model.image_size=1024', 'accelerator.run_name=step_sana_sana_sprint_0_6b_1024_variable-t_lr1e-5_step-8000_filter2_time951', 'accelerator.log_with=null', 'accelerator=deepspeed', 'optimizer=dummy', 'lr_scheduler=dummy', 'criterion.is_distributed=true', 'accelerator.deepspeed.zero_optimization.stage=3', 'accelerator.deepspeed.gradient_accumulation_steps=1', 'dataset.pseudo_preference_path=/g/data/rr81/LPO/lrm/lrm_sana/vqa_aes_clip_score_mp.csv', 'dataset.valid_split_name=validation_unique', 'dataset.test_split_name=test_unique']
Traceback (most recent call last):
  File "/g/data/rr81/LPO/lrm/lrm_sana/trainer/scripts/train.py", line 201, in main
    loss = task.train_step(model, criterion, batch)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/LPO/lrm/lrm_sana/trainer/tasks/step_sana_task.py", line 53, in train_step
    loss = criterion(model, batch)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/LPO/lrm/lrm_sana/trainer/criterions/step_clip_criterion_sana.py", line 217, in forward
    image_0_features, image_1_features, text_features = self.get_features(
                                                        ^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/LPO/lrm/lrm_sana/trainer/criterions/step_clip_criterion_sana.py", line 35, in get_features
    text_features, all_image_features = model(text_input_ids=input_ids, text_input_ids_2=input_ids_2, image_inputs=all_pixel_values, time_cond=timesteps)
                                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/amp/autocast_mode.py", line 44, in decorate_autocast
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
    ret_val = func(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/deepspeed/runtime/engine.py", line 2358, in forward
    loss = self.module(*inputs, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1881, in _call_impl
    return inner()
           ^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1829, in inner
    result = forward_call(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/LPO/lrm/lrm_sana/trainer/models/sana_preference_model.py", line 310, in forward
    image_features = self.get_image_features(
                     ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/LPO/lrm/lrm_sana/trainer/models/sana_preference_model.py", line 288, in get_image_features
    model_pred = self.transformer(**transformer_kwargs)[0]
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1881, in _call_impl
    return inner()
           ^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1829, in inner
    result = forward_call(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/diffusers/models/transformers/sana_transformer.py", line 506, in forward
    hidden_states = block(
                    ^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1881, in _call_impl
    return inner()
           ^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1829, in inner
    result = forward_call(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/diffusers/models/transformers/sana_transformer.py", line 273, in forward
    attn_output = self.attn2(
                  ^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1881, in _call_impl
    return inner()
           ^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1829, in inner
    result = forward_call(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/diffusers/models/attention_processor.py", line 605, in forward
    return self.processor(
           ^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/diffusers/models/transformers/sana_transformer.py", line 179, in __call__
    hidden_states = attn.to_out[0](hidden_states)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1881, in _call_impl
    return inner()
           ^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1818, in inner
    args_result = hook(self, args)
                  ^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 1044, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/deepspeed/runtime/zero/parameter_offload.py", line 300, in _pre_forward_module_hook
    self.pre_sub_module_forward_function(module)
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/deepspeed/runtime/zero/parameter_offload.py", line 475, in pre_sub_module_forward_function
    param_coordinator.fetch_sub_module(sub_module, forward=True)
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 1044, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
    ret_val = func(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/deepspeed/runtime/zero/partitioned_param_coordinator.py", line 325, in fetch_sub_module
    self._fetch_sub_module_impl(current_submodule, forward, is_leaf)
  File "/g/data/rr81/aev/lib/python3.11/site-packages/deepspeed/runtime/zero/partitioned_param_coordinator.py", line 377, in _fetch_sub_module_impl
    while self.__ongoing_fetch_events and self.__ongoing_fetch_events[0].query():
                                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/torch/cuda/streams.py", line 212, in query
    return super().query()
           ^^^^^^^^^^^^^^^
torch.AcceleratorError: CUDA error: misaligned address
Search for `cudaErrorMisalignedAddress' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.


Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.

Steps:   0%|          | 0/8000 [00:10<?, ?it/s]
Traceback (most recent call last):
  File "/g/data/rr81/aev/bin/accelerate", line 7, in <module>
    sys.exit(main())
             ^^^^^^
  File "/g/data/rr81/aev/lib/python3.11/site-packages/accelerate/commands/accelerate_cli.py", line 50, in main
    args.func(args)
  File "/g/data/rr81/aev/lib/python3.11/site-packages/accelerate/commands/launch.py", line 1281, in launch_command
    simple_launcher(args)
  File "/g/data/rr81/aev/lib/python3.11/site-packages/accelerate/commands/launch.py", line 869, in simple_launcher
    raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
subprocess.CalledProcessError: Command '['/g/data/rr81/aev/bin/python3.11', 'trainer/scripts/train.py', '--config-path', '/g/data/rr81/LPO/lrm/lrm_sana/trainer/conf', '--config-name', 'step_sana_base', 'accelerator.mixed_precision=BF16', 'model.model_profile=sana_sprint_0_6b_1024', 'model.pretrained_model_name_or_path=Efficient-Large-Model/Sana_Sprint_0.6B_1024px_diffusers', 'model.image_size=1024', 'accelerator.run_name=step_sana_sana_sprint_0_6b_1024_variable-t_lr1e-5_step-8000_filter2_time951', 'accelerator.log_with=null', 'accelerator=deepspeed', 'optimizer=dummy', 'lr_scheduler=dummy', 'criterion.is_distributed=true', 'accelerator.deepspeed.zero_optimization.stage=3', 'accelerator.deepspeed.gradient_accumulation_steps=1', 'dataset.pseudo_preference_path=/g/data/rr81/LPO/lrm/lrm_sana/vqa_aes_clip_score_mp.csv', 'dataset.valid_split_name=validation_unique', 'dataset.test_split_name=test_unique']' returned non-zero exit status 1.