[NPU] Avoid repeated BF16 wo_a weight transposes in DeepSeek-V4 decode (#39919)
This commit is contained in:
Jensen
@@ -0,0 +1,217 @@
|
||||
"""Compare BF16 wo_a einsum and F.linear on one Ascend NPU.
|
||||
|
||||
python test/manual/dsv4/bench_npu_wo_a.py --batch-sizes 1 8
|
||||
|
||||
Simulates one DSV4 Flash TP8 rank with one local output group. Both variants
|
||||
share the same inputs and 43 distinct layer weights by default. One iteration
|
||||
replays all --layers projections; ms_per_iteration reports their total time,
|
||||
and us_per_layer reports the average. Model loading and serving are not timed.
|
||||
F.linear is called directly, as in the model, so no optimization flags are needed.
|
||||
Optional --profile-dir exports a trace for checking weight transposes.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import statistics
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def _einsum(o, weight):
|
||||
return torch.einsum("tgd,grd->tgr", o, weight.unsqueeze(0))
|
||||
|
||||
|
||||
def _capture(fn, inputs, weights):
|
||||
def run():
|
||||
return [fn(o, w) for o, w in zip(inputs, weights)]
|
||||
|
||||
stream = torch.npu.Stream()
|
||||
stream.wait_stream(torch.npu.current_stream())
|
||||
with torch.npu.stream(stream):
|
||||
for _ in range(3):
|
||||
run()
|
||||
torch.npu.current_stream().wait_stream(stream)
|
||||
torch.npu.synchronize()
|
||||
graph = torch.npu.NPUGraph()
|
||||
with torch.npu.graph(graph, stream=stream, auto_dispatch_capture=True):
|
||||
outputs = run()
|
||||
return graph, outputs
|
||||
|
||||
|
||||
def _time_graph(graph, iterations):
|
||||
start = torch.npu.Event(enable_timing=True)
|
||||
end = torch.npu.Event(enable_timing=True)
|
||||
start.record()
|
||||
for _ in range(iterations):
|
||||
graph.replay()
|
||||
end.record()
|
||||
end.synchronize()
|
||||
return start.elapsed_time(end) / iterations
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--batch-sizes", type=int, nargs="+", default=[1, 8])
|
||||
parser.add_argument(
|
||||
"--layers",
|
||||
type=int,
|
||||
default=43,
|
||||
help="Distinct wo_a projections per graph replay.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dim", type=int, default=4096, help="Input feature width of the local group."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--o-lora-rank",
|
||||
type=int,
|
||||
default=1024,
|
||||
help="Output projection width, not a distributed rank ID.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input-stride-factor",
|
||||
type=int,
|
||||
default=8,
|
||||
help="Token stride / D; 8 models the 64-head output sliced to 8 TP-local heads.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--iterations", type=int, default=100, help="Graph replays per timing round."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--rounds", type=int, default=5, help="Timing rounds; results use the median."
|
||||
)
|
||||
parser.add_argument("--device", default="npu:0")
|
||||
parser.add_argument("--profile-dir")
|
||||
args = parser.parse_args()
|
||||
if (
|
||||
min(
|
||||
*args.batch_sizes,
|
||||
args.layers,
|
||||
args.dim,
|
||||
args.o_lora_rank,
|
||||
args.input_stride_factor,
|
||||
args.iterations,
|
||||
args.rounds,
|
||||
)
|
||||
< 1
|
||||
):
|
||||
parser.error("dimensions and iteration counts must be positive")
|
||||
|
||||
import torch_npu
|
||||
|
||||
torch.npu.set_device(args.device)
|
||||
torch.manual_seed(42)
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"torch": torch.__version__,
|
||||
"torch_npu": torch_npu.__version__,
|
||||
"device": torch.npu.get_device_name(),
|
||||
"args": vars(args),
|
||||
}
|
||||
),
|
||||
flush=True,
|
||||
)
|
||||
weights = [
|
||||
torch.randn(
|
||||
args.o_lora_rank, args.dim, device=args.device, dtype=torch.bfloat16
|
||||
)
|
||||
/ math.sqrt(args.dim)
|
||||
for _ in range(args.layers)
|
||||
]
|
||||
variants = {
|
||||
"einsum_original": _einsum,
|
||||
"linear_original": F.linear,
|
||||
}
|
||||
for tokens in args.batch_sizes:
|
||||
inputs = [
|
||||
torch.randn(
|
||||
tokens,
|
||||
args.input_stride_factor,
|
||||
args.dim,
|
||||
device=args.device,
|
||||
dtype=torch.bfloat16,
|
||||
)[:, :1, :]
|
||||
for _ in weights
|
||||
]
|
||||
graphs = {}
|
||||
outputs = {}
|
||||
for name, fn in variants.items():
|
||||
graphs[name], outputs[name] = _capture(fn, inputs, weights)
|
||||
graphs[name].replay()
|
||||
torch.npu.synchronize()
|
||||
for actual, reference in zip(
|
||||
outputs["linear_original"], outputs["einsum_original"]
|
||||
):
|
||||
torch.testing.assert_close(actual, reference, rtol=0.016, atol=0.016)
|
||||
|
||||
# A captured graph must see an in-place online weight update even when
|
||||
# the update bypasses model.post_load_weights(), as direct updates do.
|
||||
pointer = weights[0].data_ptr()
|
||||
saved = weights[0].clone()
|
||||
before_update = outputs["einsum_original"][0].clone()
|
||||
weights[0].neg_()
|
||||
for graph in graphs.values():
|
||||
graph.replay()
|
||||
torch.npu.synchronize()
|
||||
assert not torch.equal(outputs["einsum_original"][0], before_update)
|
||||
# No bias: negating W must negate the result. This also catches a stale
|
||||
# weight in the baseline graph rather than only comparing two graphs.
|
||||
expected_after_update = -before_update
|
||||
for tensors in outputs.values():
|
||||
torch.testing.assert_close(
|
||||
tensors[0], expected_after_update, rtol=0.016, atol=0.016
|
||||
)
|
||||
assert weights[0].data_ptr() == pointer
|
||||
weights[0].copy_(saved)
|
||||
torch.npu.synchronize()
|
||||
|
||||
samples = {name: [] for name in graphs}
|
||||
names = list(graphs)
|
||||
for round_id in range(args.rounds):
|
||||
# Rotate measurement order to reduce a fixed-order clock bias.
|
||||
offset = round_id % len(names)
|
||||
for name in names[offset:] + names[:offset]:
|
||||
samples[name].append(_time_graph(graphs[name], args.iterations))
|
||||
baseline = statistics.median(samples["einsum_original"])
|
||||
for name, times in samples.items():
|
||||
elapsed = statistics.median(times)
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"batch_size": tokens,
|
||||
"variant": name,
|
||||
"layers": args.layers,
|
||||
"input_stride": list(inputs[0].stride()),
|
||||
"ms_per_iteration": elapsed,
|
||||
"us_per_layer": elapsed * 1000 / args.layers,
|
||||
"speedup": baseline / elapsed,
|
||||
"samples_ms": times,
|
||||
}
|
||||
),
|
||||
flush=True,
|
||||
)
|
||||
|
||||
if args.profile_dir:
|
||||
with torch_npu.profiler.profile(
|
||||
activities=[
|
||||
torch_npu.profiler.ProfilerActivity.CPU,
|
||||
torch_npu.profiler.ProfilerActivity.NPU,
|
||||
],
|
||||
record_shapes=True,
|
||||
with_stack=False,
|
||||
on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(
|
||||
f"{args.profile_dir}/bs{tokens}"
|
||||
),
|
||||
) as profiler:
|
||||
for name, graph in graphs.items():
|
||||
with torch.autograd.profiler.record_function(name):
|
||||
graph.replay()
|
||||
torch.npu.synchronize()
|
||||
profiler.step()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user