""" Standalone script to profile the data loader performance. Isolates data loading from model training to identify bottlenecks. """ import os import random import time import subprocess import threading from collections import defaultdict from contextlib import nullcontext import numpy as np import torch from torch.utils.data import DataLoader from torch.profiler import record_function from tqdm import tqdm from data_utils import SamplingParams, BCTDataset from audioloader import AudioLoaderDataset from modules.gpt import GPTConfig, GPTTrainConfig from utils.profiling import maybe_enable_profiling # Import train.py defaults print("Loading defaults from train.py...") exec(open("train.py").read().split('exec(open("configurator.py").read())')[0]) # Override specific values for profiling data_dir = "/app2/suno/data/auk_v0" train_metas_filename = "metas_v8_tr_mini.jsonl" block_size = 8000 batch_size = 1 batch_store_size = 1 n_layer = 2 n_head = 32 d_head = 128 num_batches_to_profile = 50 num_workers_oracle = 6 # Profiling-specific settings use_torch_profiler = False # Enable torch.profiler for detailed traces torch_profiler_output_dir = "./profiler_traces" torch_profiler_wait = 2 # Wait steps before profiling torch_profiler_warmup = 2 # Warmup steps torch_profiler_active = 3 # Active profiling steps torch_profiler_repeat = 1 # Number of cycles to repeat prefetch_factor_oracle = batch_store_size * 4 device = "cuda" if torch.cuda.is_available() else "cpu" seed_offset = 1 torch.manual_seed(6006 + seed_offset) random.seed(6006 + seed_offset) np.random.seed(6006 + seed_offset) if pack: batch_size = 1 batch_size_tokens = block_size * batch_size if t_text is None: t_text = int(round(block_size * 0.1 / 256) * 256) if t_audio is None: t_audio = int(round(block_size * 0.85 / 256) * 256) class GPUMonitor: """Monitor GPU utilization in a background thread.""" def __init__(self, interval=0.1, gpu_id=None): self.interval = interval self.gpu_id = gpu_id if gpu_id is not None else 0 self.gpu_utils = [] self.gpu_mems = [] self.running = False self.thread = None def _monitor(self): while self.running: try: # Get GPU utilization and memory usage for specific GPU result = subprocess.run( [ "nvidia-smi", "--query-gpu=utilization.gpu,memory.used", "--format=csv,noheader,nounits", f"--id={self.gpu_id}", ], capture_output=True, text=True, check=True, ) # Parse output (format: "95, 12345") parts = result.stdout.strip().split(",") if len(parts) >= 2: gpu_util = float(parts[0].strip()) gpu_mem = float(parts[1].strip()) self.gpu_utils.append(gpu_util) self.gpu_mems.append(gpu_mem) except Exception as e: pass # Silently ignore errors time.sleep(self.interval) def start(self): self.running = True self.thread = threading.Thread(target=self._monitor, daemon=True) self.thread.start() def stop(self): self.running = False if self.thread: self.thread.join(timeout=1.0) def get_stats(self): if not self.gpu_utils: return None return { "mean": np.mean(self.gpu_utils), "median": np.median(self.gpu_utils), "min": np.min(self.gpu_utils), "max": np.max(self.gpu_utils), "p50": np.percentile(self.gpu_utils, 50), "p95": np.percentile(self.gpu_utils, 95), "p99": np.percentile(self.gpu_utils, 99), } def profile_dataloader(): """Profile the data loader and report detailed statistics.""" print("=" * 80) print("DATA LOADER PROFILING") print("=" * 80) print(f"Data directory: {data_dir}") print(f"Batch size: {batch_size}") print(f"Batch size tokens: {batch_size_tokens}") print(f"Batch store size: {batch_store_size}") print(f"Pack: {pack}") print(f"Oracle workers: {num_workers_oracle}") print(f"Prefetch factor: {prefetch_factor_oracle}") print(f"Number of batches: {num_batches_to_profile}") print("=" * 80) print() # Create model configs model_args = dict( n_layer=n_layer, n_head=n_head, n_kv_head=n_kv_head, d_head=d_head, block_size=block_size, bias=bias, text_vocab_size=text_vocab_size, text_codebook_size=text_codebook_size, semantic_vocab_size=semantic_vocab_size, semantic_codebook_size=semantic_codebook_size, semantic_n_codebooks=semantic_n_codebooks, semantic_rate_hz=semantic_rate_hz, semantic_shift_factor=semantic_shift_factor, semantic_type=semantic_type, t_text=t_text, t_audio=t_audio, use_text_loss=use_text_loss, use_mmbert=use_mmbert, use_vae_input=use_vae_input, output_paradigm=output_paradigm, output_distribution=output_distribution, use_hoot=use_hoot, use_ditto=use_ditto, use_rotary_pos_emb=use_rotary_pos_emb, rope_theta=rope_theta, use_qk_norm=use_qk_norm, activation_f=activation_f, embed_scale_factor=embed_scale_factor, attention_sliding_window_size=attention_sliding_window_size, global_every_n_layers=global_every_n_layers, use_continuous_semantic_input=use_continuous_semantic_input, use_block_type_embeddings=use_block_type_embeddings, repa_semantic=repa_semantic, repa_mixed_semantic=repa_mixed_semantic, repa_hoot=repa_hoot, repa_midi=repa_midi, repa_layer=repa_layer, ) train_model_args = dict(dropout=dropout, attention_type=attention_type, layer_init=layer_init) gptconf = GPTConfig(**model_args) gpttrainconf = GPTTrainConfig(**train_model_args) # Create sampling params train_sampling_params = SamplingParams( inference=False, mock_data=mock_data, suppress_text=suppress_text, mask_padding=mask_padding, pack=pack, allow_infill=allow_infill, allow_artist=allow_artist, allow_cover=allow_cover, allow_stem=allow_stem, allow_sample=allow_sample, prob_sample=prob_sample, prob_sample_from_stems=prob_sample_from_stems, sample_permutation_prob=sample_permutation_prob, max_num_audio_samples=10, allow_overpaint=allow_overpaint, allow_underpaint=allow_underpaint, allow_vox=allow_vox, prob_stem=prob_stem, allow_skip=allow_skip, allow_playlist=allow_playlist, allow_mumble=allow_mumble, allow_sfx=allow_sfx, allow_lyrics_randomize=allow_lyrics_randomize, text_loss=use_text_loss, prob_text_loss=prob_text_loss, use_mmbert=use_mmbert, prob_use_mmbert=prob_use_mmbert, use_vae_input=use_vae_input, output_paradigm=output_paradigm, output_distribution=output_distribution, noise_rng=None, use_hoot=use_hoot, use_ditto=use_ditto, max_audio_duration_s=max_audio_duration_s, n_samples_per_meta=n_samples_per_meta, use_continuous_semantic_input=use_continuous_semantic_input, use_noncausal_input=use_noncausal_input, repa_semantic=repa_semantic, repa_mixed_semantic=repa_mixed_semantic, repa_hoot=repa_hoot, repa_midi=repa_midi, prob_text_conditioning_pairs=prob_text_conditioning_pairs, prob_warp_contours=prob_warp_contours, prob_loudness_25hz=prob_loudness_25hz, prob_contour_loudness_seq=prob_contour_loudness_seq, prob_contour_spectral_centroid_seq=prob_contour_spectral_centroid_seq, prob_contour_spectral_complexity_seq=prob_contour_spectral_complexity_seq, prob_repa_hoot=prob_repa_hoot, prob_repa_midi=prob_repa_midi, prob_token_dropout=prob_token_dropout, token_dropout_pct=token_dropout_pct, dropout_codebook_pct=dropout_codebook_pct, interleave_probability=interleave_probability, ) print("Creating oracle dataset...") t_start = time.time() oracle_dataset = AudioLoaderDataset( gptconf, gpttrainconf, os.path.join(data_dir, train_metas_filename), batch_size_tokens, os.path.join(data_dir, tokenizer_filename), device, split="train", dataset_idx=None, info_path=None if train_info_filename is None else os.path.join(data_dir, train_info_filename), sampling_params=train_sampling_params, stem_active_sections_weight=stem_active_sections_weight, ) print(f"Oracle dataset created in {time.time() - t_start:.2f}s") print("Creating oracle dataloader...") oracle_dataloader = DataLoader( oracle_dataset, shuffle=False, num_workers=num_workers_oracle, prefetch_factor=prefetch_factor_oracle, batch_size=None, worker_init_fn=lambda x: random.seed(x + seed_offset * 1000), ) oracle_dataloader_iter = iter(oracle_dataloader) print("Creating BCT dataset...") dataset = BCTDataset( oracle_dataloader_iter, batch_size_tokens, gptconf, sampling_params=train_sampling_params, device=device, tokenizer_fp=os.path.join(data_dir, tokenizer_filename), save_debug_build_text=save_debug_build_text, debug_output_dir=debug_output_dir, ) print("Creating BCT dataloader...") dataloader = DataLoader( dataset, shuffle=False, num_workers=0, batch_size=None, worker_init_fn=lambda x: random.seed(x + seed_offset * 1000), ) print() print("=" * 80) print("PROFILING START") print("=" * 80) print() # Warmup print("Warming up (5 batches)...") dataloader_iter = iter(dataloader) for _ in range(5): _ = next(dataloader_iter) print("Warmup complete\n") # Profile batch_times = [] token_counts = [] seq_lengths = [] block_type_times = defaultdict(list) print(f"Profiling {num_batches_to_profile} batches...") dataloader_iter = iter(dataloader) # Start GPU monitoring # Get GPU ID from CUDA_VISIBLE_DEVICES or use current device gpu_id = int(os.environ.get("CUDA_VISIBLE_DEVICES", "0").split(",")[0]) gpu_monitor = GPUMonitor(interval=0.1, gpu_id=gpu_id) gpu_monitor.start() overall_start = time.time() # Setup torch profiler if enabled if use_torch_profiler: trace_dir = os.path.join(torch_profiler_output_dir, "traces") os.makedirs(trace_dir, exist_ok=True) print(f"Torch profiler enabled. Traces will be saved to: {trace_dir}") print( f" Wait: {torch_profiler_wait}, Warmup: {torch_profiler_warmup}, Active: {torch_profiler_active}, Repeat: {torch_profiler_repeat}" ) def trace_handler(prof): curr_trace_dir_name = "iteration_" + str(prof.step_num) curr_trace_dir = os.path.join(trace_dir, curr_trace_dir_name) if not os.path.exists(curr_trace_dir): os.makedirs(curr_trace_dir, exist_ok=True) print(f"Dumping trace at step {prof.step_num}") trace_path = f"{curr_trace_dir}/rank0_trace.json.gz" prof.export_chrome_trace(trace_path) print(f"Trace saved at: {trace_path}") from torch.profiler import profile, ProfilerActivity profiler_context = profile( activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], schedule=torch.profiler.schedule( wait=torch_profiler_wait, warmup=torch_profiler_warmup, active=torch_profiler_active, repeat=torch_profiler_repeat, ), on_trace_ready=trace_handler, record_shapes=False, profile_memory=False, with_stack=True, ) else: from contextlib import nullcontext profiler_context = nullcontext() with profiler_context as prof: for i in tqdm(range(num_batches_to_profile), desc="Profiling batches"): t_batch_start = time.time() with record_function("dataloader_next"): try: packed_sequences = next(dataloader_iter) except StopIteration: print(f"Dataloader exhausted after {i} batches") break t_batch_end = time.time() batch_time = t_batch_end - t_batch_start batch_times.append(batch_time) token_counts.append(packed_sequences.n_tokens) seq_lengths.append(packed_sequences.avg_seq_length_before_crop) # Track block types in this batch with record_function("track_block_types"): for seq in packed_sequences.block_sequences: for block in seq.blocks: block_type = str(block.spec) block_type_times[block_type].append(batch_time) if use_torch_profiler and prof: prof.step() overall_time = time.time() - overall_start # Stop GPU monitoring gpu_monitor.stop() gpu_stats = gpu_monitor.get_stats() if use_torch_profiler: trace_dir = os.path.join(torch_profiler_output_dir, "traces") print(f"\nTorch profiler traces saved to: {trace_dir}") print(f"View .json files with: chrome://tracing") print() # Calculate statistics print() print("=" * 80) print("PROFILING RESULTS") print("=" * 80) print() batch_times = np.array(batch_times) token_counts = np.array(token_counts) seq_lengths = np.array(seq_lengths) print(f"Total batches profiled: {len(batch_times)}") print(f"Overall time: {overall_time:.2f}s") print() print("BATCH TIMING:") print(f" Mean batch time: {batch_times.mean() * 1000:.1f}ms") print(f" Median batch time: {np.median(batch_times) * 1000:.1f}ms") print(f" Min batch time: {batch_times.min() * 1000:.1f}ms") print(f" Max batch time: {batch_times.max() * 1000:.1f}ms") print(f" Std batch time: {batch_times.std() * 1000:.1f}ms") print() print("THROUGHPUT:") total_tokens = token_counts.sum() tokens_per_sec = total_tokens / overall_time print(f" Total tokens: {total_tokens:,}") print(f" Tokens per second: {tokens_per_sec:,.0f}") print(f" Batches per second: {len(batch_times) / overall_time:.2f}") print() print("TOKEN STATISTICS:") print(f" Mean tokens per batch: {token_counts.mean():,.0f}") print(f" Median tokens per batch: {np.median(token_counts):,.0f}") print(f" Min tokens per batch: {token_counts.min():,}") print(f" Max tokens per batch: {token_counts.max():,}") print() print("SEQUENCE LENGTH STATISTICS:") print(f" Mean seq length: {seq_lengths.mean():,.0f}") print(f" Median seq length: {np.median(seq_lengths):,.0f}") print(f" Min seq length: {seq_lengths.min():,}") print(f" Max seq length: {seq_lengths.max():,}") print() # Percentiles print("BATCH TIME PERCENTILES:") for p in [50, 75, 90, 95, 99]: print(f" P{p}: {np.percentile(batch_times, p) * 1000:.1f}ms") print() # Detect outliers threshold = batch_times.mean() + 2 * batch_times.std() outliers = batch_times > threshold if outliers.any(): print(f"OUTLIERS (> mean + 2*std):") print(f" Number of outliers: {outliers.sum()}") print(f" Outlier percentage: {outliers.sum() / len(batch_times) * 100:.1f}%") print(f" Mean outlier time: {batch_times[outliers].mean() * 1000:.1f}ms") print() # Memory usage if available if torch.cuda.is_available(): print("GPU MEMORY:") print(f" Allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB") print(f" Reserved: {torch.cuda.memory_reserved() / 1e9:.2f} GB") print() # GPU utilization stats if gpu_stats: print("GPU UTILIZATION:") print(f" Mean: {gpu_stats['mean']:.1f}%") print(f" Median: {gpu_stats['median']:.1f}%") print(f" Min: {gpu_stats['min']:.1f}%") print(f" Max: {gpu_stats['max']:.1f}%") print(f" P95: {gpu_stats['p95']:.1f}%") print(f" P99: {gpu_stats['p99']:.1f}%") print() # Block type analysis if block_type_times: print("BLOCK TYPE TIMING ANALYSIS:") print("(Average batch time when block type is present)") print() # Sort by average time (descending) block_stats = [] for block_type, times in block_type_times.items(): block_stats.append( { "type": block_type, "count": len(times), "avg_time": np.mean(times) * 1000, "median_time": np.median(times) * 1000, "min_time": np.min(times) * 1000, "max_time": np.max(times) * 1000, } ) block_stats.sort(key=lambda x: x["avg_time"], reverse=True) print(f"{'Block Type':<30} {'Count':<8} {'Avg Time':<12} {'Median':<12} {'Min':<12} {'Max':<12}") print("-" * 90) for stat in block_stats: print( f"{stat['type']:<30} {stat['count']:<8} " f"{stat['avg_time']:>10.1f}ms {stat['median_time']:>10.1f}ms " f"{stat['min_time']:>10.1f}ms {stat['max_time']:>10.1f}ms" ) print() print("=" * 80) print("PROFILING COMPLETE") print("=" * 80) if __name__ == "__main__": import sys # Get config keys for override from command line config_keys = [ k for k, v in globals().items() if not k.startswith("_") and isinstance(v, (int, float, bool, str, type(None))) ] # Run configurator if it exists (allows command-line overrides) config_file = "configurator.py" if os.path.exists(config_file): print(f"Loading configuration from {config_file}") exec(open(config_file).read()) # Print final config print("\nFinal Configuration:") print(f" data_dir: {data_dir}") print(f" train_metas_filename: {train_metas_filename}") print(f" block_size: {block_size}") print(f" batch_size: {batch_size}") print(f" batch_store_size: {batch_store_size}") print(f" n_layer: {n_layer}") print(f" n_head: {n_head}") print(f" d_head: {d_head}") print() profile_dataloader()