File size: 5,796 Bytes
51ce5b3 | 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 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | #!/usr/bin/env python
"""Measure the TinyVLA training step on real hardware, synthetic data.
Builds the model out of the raw modules (SemanticPath / FlowMatchingExpert /
projections) so lerobot isn't needed, and reproduces _conditioning + the flow
loss exactly as modeling_tinyvla.py does. Sweeps: precision, dedup of the cam0
vision-tower pass, torch.compile, batch size.
Prints samples/s -> hours per 1M frames and per the 8.64M-sample C-scaled budget.
"""
from __future__ import annotations
import argparse
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
from tinyvla.modules.expert import FlowMatchingExpert
from tinyvla.modules.semantic import SemanticPath
class Bench(nn.Module):
"""_conditioning + forward of TinyVLAPolicy, minus the lerobot wrapper."""
def __init__(self, d=512, max_state=64, max_action=64, chunk=50,
lm_layers=12, readout=8, freeze_lm=False, freeze_vision=True,
model_name="Qwen/Qwen3.5-0.8B"):
super().__init__()
self.semantic = SemanticPath(
model_name=model_name, num_layers=lm_layers, num_readout=readout,
out_dim=d, image_size=256, freeze_lm=freeze_lm, freeze_vision=freeze_vision,
)
self.spatial_proj = nn.Linear(self.semantic.visual_hidden_size, d)
self.camera_emb = nn.Embedding(3, d)
self.state_proj = nn.Linear(max_state, d)
self.embodiment_emb = nn.Embedding(16, d)
self.expert = FlowMatchingExpert(action_dim=max_action, chunk_size=chunk, d_model=d)
self.chunk, self.max_action = chunk, max_action
def forward(self, batch, dedup=True):
cam0, cam1 = batch["cam0"], batch["cam1"]
f0 = self.semantic.encode_image(cam0)
f1 = self.semantic.encode_image(cam1)
if dedup:
latent = self.semantic(None, batch["tok"], batch["mask"], image_embeds=f0)
else:
latent = self.semantic(cam0, batch["tok"], batch["mask"])
spatial = torch.cat([
self.spatial_proj(f0) + self.camera_emb.weight[0][None, None],
self.spatial_proj(f1) + self.camera_emb.weight[1][None, None],
], dim=1)
state_tok = self.state_proj(batch["state"])[:, None]
emb = self.embodiment_emb(batch["emb_id"])[:, None]
cond = torch.cat([latent, spatial, state_tok, emb], dim=1)
actions = batch["actions"]
b = actions.shape[0]
t = torch.rand(b, device=actions.device) * 0.999 + 0.001
noise = torch.randn_like(actions)
x_t = t[:, None, None] * noise + (1 - t[:, None, None]) * actions
pred = self.expert(x_t, t, cond)
return F.mse_loss(pred, noise - actions)
def make_batch(b, dev, max_state=64, max_action=64, chunk=50, lm_len=48):
return {
"cam0": torch.rand(b, 3, 256, 256, device=dev),
"cam1": torch.rand(b, 3, 256, 256, device=dev),
"tok": torch.randint(1000, 5000, (b, lm_len), device=dev),
"mask": torch.ones(b, lm_len, dtype=torch.bool, device=dev),
"state": torch.randn(b, max_state, device=dev),
"actions": torch.randn(b, chunk, max_action, device=dev),
"emb_id": torch.randint(0, 16, (b,), device=dev),
}
def run(model, opt, batch, dedup, steps, warmup, autocast):
for i in range(warmup + steps):
if i == warmup:
torch.cuda.synchronize()
t0 = time.time()
opt.zero_grad(set_to_none=True)
with torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast):
loss = model(batch, dedup=dedup)
loss.backward()
opt.step()
torch.cuda.synchronize()
return (time.time() - t0) / steps
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model-name", default="Qwen/Qwen3.5-0.8B")
ap.add_argument("--steps", type=int, default=30)
ap.add_argument("--warmup", type=int, default=10)
ap.add_argument("--batches", type=int, nargs="+", default=[48, 96, 144, 192])
ap.add_argument("--compile", action="store_true")
ap.add_argument("--freeze-lm", action="store_true")
args = ap.parse_args()
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
dev = torch.device("cuda")
model = Bench(freeze_lm=args.freeze_lm, model_name=args.model_name).to(dev)
n_train = sum(p.numel() for p in model.parameters() if p.requires_grad)
n_all = sum(p.numel() for p in model.parameters())
print(f"params: {n_all/1e6:.1f}M total, {n_train/1e6:.1f}M trainable "
f"(freeze_lm={args.freeze_lm}, vision frozen)")
opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad],
lr=1e-4, betas=(0.9, 0.95), weight_decay=1e-10, fused=True)
if args.compile:
model.expert = torch.compile(model.expert, dynamic=False)
BUDGET = 8_640_000 # C-scaled: 60k steps x eff batch 144
print(f"\n{'batch':>6} {'dedup':>6} {'ms/step':>9} {'samples/s':>10} "
f"{'ч на 1M кадров':>15} {'ч на 8.64M':>11} {'VRAM GB':>8}")
for b in args.batches:
batch = make_batch(b, dev)
for dedup in (True, False):
torch.cuda.reset_peak_memory_stats()
try:
dt = run(model, opt, batch, dedup, args.steps, args.warmup, autocast=True)
except torch.cuda.OutOfMemoryError:
print(f"{b:>6} {str(dedup):>6} OOM")
torch.cuda.empty_cache()
continue
sps = b / dt
mem = torch.cuda.max_memory_allocated() / 1e9
print(f"{b:>6} {str(dedup):>6} {dt*1000:>9.1f} {sps:>10.0f} "
f"{1e6/sps/3600:>15.2f} {BUDGET/sps/3600:>11.2f} {mem:>8.1f}", flush=True)
if __name__ == "__main__":
main()
|