101 lines
4.2 KiB
Python
101 lines
4.2 KiB
Python
import json
|
|
import os
|
|
import threading
|
|
import time
|
|
|
|
import torch
|
|
|
|
from prain import profile
|
|
from tests.helpers import TmpCase
|
|
|
|
|
|
class Inner(torch.nn.Module):
|
|
def forward(self, x):
|
|
time.sleep(0.002)
|
|
return x
|
|
|
|
|
|
class Outer(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.layers = torch.nn.ModuleList([Inner(), Inner()])
|
|
self.head = Inner()
|
|
|
|
def forward(self, x):
|
|
for layer in self.layers:
|
|
x = layer(x)
|
|
with profile.span("copy"):
|
|
time.sleep(0.001)
|
|
return self.head(x)
|
|
|
|
|
|
class TestProfile(TmpCase):
|
|
def setUp(self):
|
|
super().setUp()
|
|
self.syncs = []
|
|
saved = profile.DIR, profile.SYNC[0]
|
|
profile.DIR, profile.SYNC[0] = self.dir, lambda: self.syncs.append(1)
|
|
self.addCleanup(lambda: (setattr(profile, "DIR", saved[0]),
|
|
profile.SYNC.__setitem__(0, saved[1])))
|
|
|
|
def steps(self, model, tokens):
|
|
for n in tokens:
|
|
profile.step_begin()
|
|
with profile.span("execute_model"):
|
|
profile.note_step(model, n)
|
|
model(torch.zeros(1))
|
|
profile.step_end()
|
|
|
|
def folded(self, kind):
|
|
with open(os.path.join(self.dir, f"{kind}.folded")) as f:
|
|
return {stack: int(us) for stack, us in (l.rsplit(" ", 1) for l in f.read().splitlines())}
|
|
|
|
def test_records_the_requested_steps_as_folded_stacks(self):
|
|
model = Outer()
|
|
self.steps(model, [1]) # no start file: nothing recorded, nothing synced
|
|
self.assertEqual(self.syncs, [])
|
|
self.write("start", "3")
|
|
self.steps(model, [1, 1, 300, 1]) # records 2 decode steps and 1 prefill, not the 4th
|
|
self.assertFalse(os.path.exists(os.path.join(self.dir, "start")))
|
|
with open(os.path.join(self.dir, "steps.json")) as f:
|
|
self.assertEqual(json.load(f)["steps"], {"decode": 2, "prefill": 1})
|
|
decode = self.folded("decode")
|
|
self.assertEqual(set(decode), {"step", "step;execute_model", "step;execute_model;Outer",
|
|
"step;execute_model;Outer;Inner", # both layers, merged
|
|
"step;execute_model;Outer;copy",
|
|
"step;execute_model;Outer;head:Inner"})
|
|
for stack, ms in (("Outer;Inner", 8), ("Outer;head:Inner", 4), ("Outer;copy", 2)):
|
|
self.assertGreaterEqual(decode[f"step;execute_model;{stack}"], ms * 1000)
|
|
self.assertLess(decode[f"step;execute_model;{stack}"], ms * 1000 * 3)
|
|
self.assertGreater(self.folded("prefill")["step;execute_model;Outer;Inner"], 4000)
|
|
with open(os.path.join(self.dir, "decode.calls")) as f:
|
|
calls = dict(line.rsplit(" ", 1) for line in f.read().splitlines())
|
|
self.assertEqual(calls["step;execute_model;Outer;Inner"], "4") # 2 layers x 2 steps
|
|
self.assertEqual(calls["step"], "2")
|
|
recorded = len(self.syncs)
|
|
self.steps(model, [1]) # finished: the module hooks are gone
|
|
self.assertEqual(len(self.syncs), recorded)
|
|
self.assertFalse(model.head._forward_hooks or model.head._forward_pre_hooks)
|
|
|
|
def test_hooks_leave_other_threads_inputs_and_outputs_alone(self):
|
|
model = Outer()
|
|
self.write("start", "1")
|
|
profile.step_begin()
|
|
profile.note_step(model, 1) # installs the hooks
|
|
out = []
|
|
thread = threading.Thread(target=lambda: out.append(model(torch.ones(1))))
|
|
thread.start()
|
|
thread.join()
|
|
self.assertTrue(torch.equal(out[0], torch.ones(1)))
|
|
profile.step_end()
|
|
|
|
def test_svg_shows_time_per_step(self):
|
|
lines = ["step 500", "step;attn 3000", "step;attn;copy 1000"]
|
|
svg = profile.svg(lines, steps=2, title="decode")
|
|
self.assertTrue(svg.startswith("<svg"))
|
|
self.assertIn("attn: 2.0 ms per step (88.9%)", svg) # 4000 us of 4500, over 2 steps
|
|
self.assertIn("copy: 0.5 ms per step", svg)
|
|
self.assertEqual(profile.top(lines, steps=2)[0], ("attn", 1.5, 66.7, None)) # self time
|
|
calls = ["step;attn 10", "step;attn;copy 4", "step 2"]
|
|
self.assertEqual(profile.top(lines, steps=2, calls=calls)[0], ("attn", 1.5, 66.7, 5.0))
|