measure julia flux mlp forward criterion
Co-Authored-By: Codex <codex@openai.com>
This commit is contained in:
Codex
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,8 @@
|
||||
name = "JuliaFluxMLP"
|
||||
uuid = "9fa932db-d4ab-4cbf-935e-32425028f0bc"
|
||||
version = "0.1.0"
|
||||
|
||||
[deps]
|
||||
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
|
||||
Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
|
||||
JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
|
||||
@@ -0,0 +1,45 @@
|
||||
# Julia Flux MLP Forward Benchmark
|
||||
|
||||
Criterion 4 for the Julia port evaluation.
|
||||
|
||||
Compares PyTorch `DeepCFRMLP` against a Flux/CUDA implementation with the
|
||||
same shape and identical exported weights:
|
||||
|
||||
- input dim: 365
|
||||
- hidden size: 512
|
||||
- hidden layers: 3
|
||||
- output dim: 22
|
||||
- activation: ReLU
|
||||
|
||||
Run:
|
||||
|
||||
```bash
|
||||
uv run python experiments/julia_flux_mlp/bench_pytorch_vs_flux.py
|
||||
```
|
||||
|
||||
The runner exports PyTorch weights and inputs to `runs/tmp/`, runs Flux on the
|
||||
same payload, checks output equivalence, then writes `results.json`.
|
||||
|
||||
## Results (2026-05-07)
|
||||
|
||||
Host GPU: NVIDIA GeForce RTX 3090. Timing uses 10 warmup forwards, then 100
|
||||
timed forwards, with CUDA synchronized around the timed loop in both runtimes.
|
||||
|
||||
| Backend | batch | forward ms | μs/state | ratio vs PyTorch |
|
||||
| --- | ---: | ---: | ---: | ---: |
|
||||
| PyTorch | 1 | 0.0829 | 82.8755 | 1.00× |
|
||||
| Flux | 1 | 0.1669 | 166.8867 | 2.01× |
|
||||
| PyTorch | 64 | 0.0927 | 1.4477 | 1.00× |
|
||||
| Flux | 64 | 0.1909 | 2.9836 | 2.06× |
|
||||
| PyTorch | 256 | 0.0877 | 0.3427 | 1.00× |
|
||||
| Flux | 256 | 0.1758 | 0.6867 | 2.00× |
|
||||
|
||||
Maximum output difference: `5.215e-08`.
|
||||
|
||||
Criterion 4 threshold:
|
||||
|
||||
- bs=64 must be within ±20% of PyTorch.
|
||||
- bs=1 and bs=256 must be within ±30% of PyTorch.
|
||||
|
||||
**Verdict:** FAIL. Flux/CUDA is ~2.0× slower than PyTorch at all measured
|
||||
batch sizes for this model shape.
|
||||
@@ -0,0 +1,120 @@
|
||||
using CUDA
|
||||
using Flux
|
||||
using JSON
|
||||
using Printf
|
||||
|
||||
const WARMUP_ITERS = 10
|
||||
const TIMED_ITERS = 100
|
||||
|
||||
function dense_from_payload(layer)
|
||||
weight = Float32.(layer["weight"])
|
||||
bias = Float32.(layer["bias"])
|
||||
out_dim = Int(layer["out_dim"])
|
||||
in_dim = Int(layer["in_dim"])
|
||||
dense = Dense(in_dim => out_dim)
|
||||
dense.weight .= reshape(weight, in_dim, out_dim)'
|
||||
dense.bias .= bias
|
||||
return dense
|
||||
end
|
||||
|
||||
function build_model(payload)
|
||||
layers = Any[]
|
||||
dense_layers = payload["layers"]
|
||||
for (index, layer) in enumerate(dense_layers)
|
||||
push!(layers, dense_from_payload(layer))
|
||||
if index < length(dense_layers)
|
||||
push!(layers, relu)
|
||||
end
|
||||
end
|
||||
return Flux.fmap(cu, Chain(layers...))
|
||||
end
|
||||
|
||||
function input_matrix(batch_payload)
|
||||
batch_size = Int(batch_payload["batch_size"])
|
||||
input_dim = Int(batch_payload["input_dim"])
|
||||
flat = Float32.(batch_payload["input"])
|
||||
x = Matrix{Float32}(undef, input_dim, batch_size)
|
||||
@inbounds for sample in 1:batch_size
|
||||
source_base = (sample - 1) * input_dim
|
||||
for feature in 1:input_dim
|
||||
x[feature, sample] = flat[source_base + feature]
|
||||
end
|
||||
end
|
||||
return cu(x)
|
||||
end
|
||||
|
||||
function flatten_output(y)
|
||||
cpu = Array(y)
|
||||
out_dim, batch_size = size(cpu)
|
||||
flat = Vector{Float32}(undef, batch_size * out_dim)
|
||||
@inbounds for sample in 1:batch_size
|
||||
dest_base = (sample - 1) * out_dim
|
||||
for output in 1:out_dim
|
||||
flat[dest_base + output] = cpu[output, sample]
|
||||
end
|
||||
end
|
||||
return flat
|
||||
end
|
||||
|
||||
function timed_forward_ms(model, x)
|
||||
CUDA.synchronize()
|
||||
for _ in 1:WARMUP_ITERS
|
||||
model(x)
|
||||
end
|
||||
CUDA.synchronize()
|
||||
|
||||
elapsed = @elapsed begin
|
||||
for _ in 1:TIMED_ITERS
|
||||
model(x)
|
||||
end
|
||||
CUDA.synchronize()
|
||||
end
|
||||
return elapsed * 1000.0 / TIMED_ITERS
|
||||
end
|
||||
|
||||
function benchmark(payload)
|
||||
CUDA.allowscalar(false)
|
||||
model = build_model(payload)
|
||||
results = Dict{String,Any}()
|
||||
max_abs_diff = 0.0
|
||||
for batch_payload in payload["batches"]
|
||||
batch_size = Int(batch_payload["batch_size"])
|
||||
x = input_matrix(batch_payload)
|
||||
y = model(x)
|
||||
CUDA.synchronize()
|
||||
actual = flatten_output(y)
|
||||
expected = Float32.(batch_payload["expected_output"])
|
||||
diff = maximum(abs.(actual .- expected))
|
||||
max_abs_diff = max(max_abs_diff, Float64(diff))
|
||||
ms = timed_forward_ms(model, x)
|
||||
results[string(batch_size)] = Dict(
|
||||
"forward_ms" => ms,
|
||||
"us_per_state" => ms * 1000.0 / batch_size,
|
||||
"max_abs_diff" => Float64(diff),
|
||||
)
|
||||
end
|
||||
return Dict(
|
||||
"lang" => "Julia/Flux",
|
||||
"device" => string(CUDA.name(CUDA.device())),
|
||||
"timed_iters" => TIMED_ITERS,
|
||||
"warmup_iters" => WARMUP_ITERS,
|
||||
"max_abs_diff" => max_abs_diff,
|
||||
"batches" => results,
|
||||
)
|
||||
end
|
||||
|
||||
function print_json(result)
|
||||
println(JSON.json(result))
|
||||
end
|
||||
|
||||
function main()
|
||||
if length(ARGS) != 1
|
||||
println(stderr, "usage: julia --project=experiments/julia_flux_mlp experiments/julia_flux_mlp/bench_flux.jl <payload.json>")
|
||||
exit(2)
|
||||
end
|
||||
payload = JSON.parsefile(ARGS[1])
|
||||
result = benchmark(payload)
|
||||
print_json(result)
|
||||
end
|
||||
|
||||
main()
|
||||
@@ -0,0 +1,188 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import subprocess
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from coolrl_lost_cities.games.classic.deep_cfr.networks import DeepCFRMLP
|
||||
|
||||
ROOT = Path(__file__).resolve().parent
|
||||
REPO_ROOT = ROOT.parents[1]
|
||||
BATCH_SIZES = (1, 64, 256)
|
||||
INPUT_DIM = 365
|
||||
HIDDEN_SIZE = 512
|
||||
NUM_LAYERS = 3
|
||||
OUTPUT_DIM = 22
|
||||
WARMUP_ITERS = 10
|
||||
TIMED_ITERS = 100
|
||||
SEED = 20260507
|
||||
MAX_OUTPUT_DIFF = 1e-4
|
||||
|
||||
|
||||
def _julia_executable() -> str:
|
||||
local = REPO_ROOT / "tools" / "julia" / "current" / "bin" / "julia"
|
||||
if local.exists():
|
||||
return str(local)
|
||||
return "julia"
|
||||
|
||||
|
||||
def _linear_layers(model: DeepCFRMLP) -> list[torch.nn.Linear]:
|
||||
return [layer for layer in model.net if isinstance(layer, torch.nn.Linear)]
|
||||
|
||||
|
||||
def _export_layers(model: DeepCFRMLP) -> list[dict[str, Any]]:
|
||||
layers = []
|
||||
for layer in _linear_layers(model):
|
||||
weight = layer.weight.detach().cpu().contiguous()
|
||||
bias = layer.bias.detach().cpu().contiguous()
|
||||
layers.append(
|
||||
{
|
||||
"in_dim": int(weight.shape[1]),
|
||||
"out_dim": int(weight.shape[0]),
|
||||
"weight": weight.flatten().tolist(),
|
||||
"bias": bias.flatten().tolist(),
|
||||
}
|
||||
)
|
||||
return layers
|
||||
|
||||
|
||||
def _time_pytorch(model: DeepCFRMLP, x: torch.Tensor) -> float:
|
||||
with torch.inference_mode():
|
||||
for _ in range(WARMUP_ITERS):
|
||||
model(x)
|
||||
torch.cuda.synchronize()
|
||||
start = time.perf_counter()
|
||||
for _ in range(TIMED_ITERS):
|
||||
model(x)
|
||||
torch.cuda.synchronize()
|
||||
return float((time.perf_counter() - start) * 1000.0 / TIMED_ITERS)
|
||||
|
||||
|
||||
def _build_payload_and_pytorch_results(payload_path: Path) -> dict[str, Any]:
|
||||
if not torch.cuda.is_available():
|
||||
raise SystemExit("CUDA is not available to PyTorch; aborting Criterion 4")
|
||||
torch.manual_seed(SEED)
|
||||
torch.cuda.manual_seed_all(SEED)
|
||||
torch.backends.cuda.matmul.allow_tf32 = False
|
||||
torch.backends.cudnn.allow_tf32 = False
|
||||
|
||||
model = (
|
||||
DeepCFRMLP(
|
||||
INPUT_DIM,
|
||||
OUTPUT_DIM,
|
||||
HIDDEN_SIZE,
|
||||
num_layers=NUM_LAYERS,
|
||||
activation="relu",
|
||||
)
|
||||
.eval()
|
||||
.cuda()
|
||||
)
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"input_dim": INPUT_DIM,
|
||||
"hidden_size": HIDDEN_SIZE,
|
||||
"num_layers": NUM_LAYERS,
|
||||
"output_dim": OUTPUT_DIM,
|
||||
"activation": "relu",
|
||||
"layers": _export_layers(model),
|
||||
"batches": [],
|
||||
}
|
||||
pytorch_batches: dict[str, Any] = {}
|
||||
generator = torch.Generator(device="cpu").manual_seed(SEED + 1)
|
||||
with torch.inference_mode():
|
||||
for batch_size in BATCH_SIZES:
|
||||
x_cpu = torch.randn(batch_size, INPUT_DIM, generator=generator, dtype=torch.float32)
|
||||
x_gpu = x_cpu.cuda()
|
||||
y_cpu = model(x_gpu).detach().cpu().contiguous()
|
||||
forward_ms = _time_pytorch(model, x_gpu)
|
||||
payload["batches"].append(
|
||||
{
|
||||
"batch_size": batch_size,
|
||||
"input_dim": INPUT_DIM,
|
||||
"input": x_cpu.contiguous().flatten().tolist(),
|
||||
"expected_output": y_cpu.flatten().tolist(),
|
||||
}
|
||||
)
|
||||
pytorch_batches[str(batch_size)] = {
|
||||
"forward_ms": forward_ms,
|
||||
"us_per_state": forward_ms * 1000.0 / batch_size,
|
||||
}
|
||||
|
||||
payload_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
payload_path.write_text(json.dumps(payload), encoding="utf-8")
|
||||
return {
|
||||
"lang": "PyTorch",
|
||||
"device": torch.cuda.get_device_name(0),
|
||||
"timed_iters": TIMED_ITERS,
|
||||
"warmup_iters": WARMUP_ITERS,
|
||||
"batches": pytorch_batches,
|
||||
}
|
||||
|
||||
|
||||
def _run_flux(payload_path: Path) -> dict[str, Any]:
|
||||
proc = subprocess.run(
|
||||
[
|
||||
_julia_executable(),
|
||||
"--project=experiments/julia_flux_mlp",
|
||||
"experiments/julia_flux_mlp/bench_flux.jl",
|
||||
str(payload_path),
|
||||
],
|
||||
cwd=REPO_ROOT,
|
||||
check=True,
|
||||
text=True,
|
||||
capture_output=True,
|
||||
)
|
||||
return json.loads(proc.stdout)
|
||||
|
||||
|
||||
def _verdict(pytorch: dict[str, Any], flux: dict[str, Any]) -> tuple[str, dict[str, float]]:
|
||||
ratios = {}
|
||||
for batch_size in BATCH_SIZES:
|
||||
key = str(batch_size)
|
||||
ratio = flux["batches"][key]["us_per_state"] / pytorch["batches"][key]["us_per_state"]
|
||||
ratios[key] = ratio
|
||||
within_64 = 0.8 <= ratios["64"] <= 1.2
|
||||
within_1 = 0.7 <= ratios["1"] <= 1.3
|
||||
within_256 = 0.7 <= ratios["256"] <= 1.3
|
||||
return ("PASS" if within_64 and within_1 and within_256 else "FAIL"), ratios
|
||||
|
||||
|
||||
def _print_table(
|
||||
pytorch: dict[str, Any], flux: dict[str, Any], ratios: dict[str, float], verdict: str
|
||||
) -> None:
|
||||
print("Backend batch forward_ms us/state ratio vs PyTorch")
|
||||
for batch_size in BATCH_SIZES:
|
||||
key = str(batch_size)
|
||||
p = pytorch["batches"][key]
|
||||
f = flux["batches"][key]
|
||||
print(
|
||||
f"PyTorch {batch_size:5d} {p['forward_ms']:10.4f} {p['us_per_state']:8.4f} 1.00x"
|
||||
)
|
||||
print(
|
||||
f"Flux {batch_size:5d} {f['forward_ms']:10.4f} "
|
||||
f"{f['us_per_state']:8.4f} {ratios[key]:.2f}x"
|
||||
)
|
||||
print(f"max output diff: {flux['max_abs_diff']:.3e}")
|
||||
print(f"Criterion 4 verdict: {verdict}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
payload_path = REPO_ROOT / "runs" / "tmp" / "julia_flux_mlp_payload.json"
|
||||
pytorch = _build_payload_and_pytorch_results(payload_path)
|
||||
flux = _run_flux(payload_path)
|
||||
if float(flux["max_abs_diff"]) > MAX_OUTPUT_DIFF:
|
||||
raise SystemExit(f"Flux output mismatch: max_abs_diff={flux['max_abs_diff']:.3e}")
|
||||
verdict, ratios = _verdict(pytorch, flux)
|
||||
result = {"pytorch": pytorch, "flux": flux, "ratios": ratios, "verdict": verdict}
|
||||
result_path = ROOT / "results.json"
|
||||
result_path.write_text(json.dumps(result, indent=2, sort_keys=True), encoding="utf-8")
|
||||
_print_table(pytorch, flux, ratios, verdict)
|
||||
print(f"wrote {result_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,52 @@
|
||||
{
|
||||
"flux": {
|
||||
"batches": {
|
||||
"1": {
|
||||
"forward_ms": 0.16688673,
|
||||
"max_abs_diff": 1.4901161193847656e-08,
|
||||
"us_per_state": 166.88673
|
||||
},
|
||||
"256": {
|
||||
"forward_ms": 0.17579791,
|
||||
"max_abs_diff": 2.9802322387695312e-08,
|
||||
"us_per_state": 0.6867105859375
|
||||
},
|
||||
"64": {
|
||||
"forward_ms": 0.19094904,
|
||||
"max_abs_diff": 5.21540641784668e-08,
|
||||
"us_per_state": 2.98357875
|
||||
}
|
||||
},
|
||||
"device": "NVIDIA GeForce RTX 3090",
|
||||
"lang": "Julia/Flux",
|
||||
"max_abs_diff": 5.21540641784668e-08,
|
||||
"timed_iters": 100,
|
||||
"warmup_iters": 10
|
||||
},
|
||||
"pytorch": {
|
||||
"batches": {
|
||||
"1": {
|
||||
"forward_ms": 0.0828754500253126,
|
||||
"us_per_state": 82.8754500253126
|
||||
},
|
||||
"256": {
|
||||
"forward_ms": 0.0877418098389171,
|
||||
"us_per_state": 0.34274144468326995
|
||||
},
|
||||
"64": {
|
||||
"forward_ms": 0.09265014989068732,
|
||||
"us_per_state": 1.4476585920419893
|
||||
}
|
||||
},
|
||||
"device": "NVIDIA GeForce RTX 3090",
|
||||
"lang": "PyTorch",
|
||||
"timed_iters": 100,
|
||||
"warmup_iters": 10
|
||||
},
|
||||
"ratios": {
|
||||
"1": 2.0137052643337428,
|
||||
"256": 2.0035819904187386,
|
||||
"64": 2.0609684951971476
|
||||
},
|
||||
"verdict": "FAIL"
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
module JuliaFluxMLP
|
||||
|
||||
end
|
||||
Reference in New Issue
Block a user