File size: 4,057 Bytes
5dc80b3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | #!/usr/bin/env python3
"""
NVIDIA Nsight Systems Profiler for HRM Memory Tiering.
This script runs a few iterations of both the baseline and tiered models
and is designed to be executed via `nsys profile`.
Usage:
nsys profile -t cuda,nvtx --stats=true --force-overwrite=true -o hrm_profile python run_nsys_profiler.py
"""
import torch
import torch.cuda.nvtx as nvtx
import argparse
from models.hrm.hrm_act_v1 import HierarchicalReasoningModel_ACTV1
from models.hrm.hrm_tiered import HRM_Tiered
from models.memory_tier import MemoryTierManager
from run_training_comparison import DummyLossModel, create_dummy_batch
def profile_model(model_name, model, batch, iterations, device):
print(f"Profiling {model_name}...")
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
model.train()
# Warmup
for _ in range(2):
optimizer.zero_grad()
carry = model.initial_carry(batch)
carry.inner_carry.z_H = carry.inner_carry.z_H.to(device)
carry.inner_carry.z_L = carry.inner_carry.z_L.to(device)
carry.steps = carry.steps.to(device)
carry.halted = carry.halted.to(device)
carry.current_data = {k: v.to(device) for k, v in carry.current_data.items()}
_, loss, _, _, _ = model(carry, batch, return_keys=[])
loss.backward()
optimizer.step()
torch.cuda.synchronize()
# Profiling Phase
with torch.autograd.profiler.emit_nvtx():
nvtx.range_push(f"{model_name}_Training_Loop")
for i in range(iterations):
nvtx.range_push(f"Iteration_{i}")
optimizer.zero_grad()
nvtx.range_push("Forward_Pass")
carry = model.initial_carry(batch)
carry.inner_carry.z_H = carry.inner_carry.z_H.to(device)
carry.inner_carry.z_L = carry.inner_carry.z_L.to(device)
carry.steps = carry.steps.to(device)
carry.halted = carry.halted.to(device)
carry.current_data = {k: v.to(device) for k, v in carry.current_data.items()}
_, loss, _, _, _ = model(carry, batch, return_keys=[])
nvtx.range_pop() # End Forward
nvtx.range_push("Backward_Pass")
loss.backward()
optimizer.step()
nvtx.range_pop() # End Backward
nvtx.range_pop() # End Iteration
nvtx.range_pop() # End Loop
torch.cuda.synchronize()
print(f"Finished {model_name}.\n")
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--batch-size', type=int, default=16)
parser.add_argument('--seq-len', type=int, default=128)
parser.add_argument('--hidden-size', type=int, default=1024)
parser.add_argument('--iterations', type=int, default=5)
args = parser.parse_args()
device = torch.device('cuda')
vocab_size = 32
config_dict = {
'batch_size': args.batch_size,
'seq_len': args.seq_len,
'puzzle_emb_ndim': 0,
'num_puzzle_identifiers': args.batch_size,
'vocab_size': vocab_size,
'H_cycles': 2,
'L_cycles': 2,
'H_layers': 4,
'L_layers': 4,
'hidden_size': args.hidden_size,
'expansion': 4.0,
'num_heads': 8,
'pos_encodings': 'rope',
'halt_max_steps': 1,
'halt_exploration_prob': 0.0,
}
batch = create_dummy_batch(args.batch_size, args.seq_len, vocab_size, device)
# 1. Baseline
baseline = HierarchicalReasoningModel_ACTV1(config_dict).to(device)
baseline_wrapped = DummyLossModel(baseline)
profile_model("HRM_Baseline", baseline_wrapped, batch, args.iterations, device)
del baseline_wrapped, baseline
torch.cuda.empty_cache()
# 2. Tiered
mem_mgr = MemoryTierManager(device=device, enable_tracking=False)
tiered = HRM_Tiered(config_dict, memory_manager=mem_mgr).to(device)
tiered_wrapped = DummyLossModel(tiered)
profile_model("HRM_Tiered", tiered_wrapped, batch, args.iterations, device)
if __name__ == "__main__":
main()
|