Add 0.5B Alpamayo validation

2026-07-11 13:08:42 +08:00
commit fb1bd51046
34 changed files with 6125 additions and 0 deletions
@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+from dataclasses import asdict
+
+import torch
+
+from alpamayo_lite.config import load_config
+from alpamayo_lite.data import load_manifest
+from alpamayo_lite.expert import ExpertArchitecture, KVCachedDiffusionExpert
+from alpamayo_lite.prompts import build_expert_prefix, build_vlm_prefix
+from alpamayo_lite.runtime import (
+    empty_device_cache,
+    parameter_counts,
+    resolve_dtype,
+    seed_everything,
+    select_device,
+    write_json,
+)
+from alpamayo_lite.vlm import (
+    configure_stage1_trainability,
+    extract_prefix_kv,
+    generate_reasoning,
+    load_vlm,
+    vlm_text_config,
+)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="configs/local_overfit.yaml")
+    args = parser.parse_args()
+    cfg = load_config(args.config)
+    seed_everything(cfg.seed)
+    device = select_device()
+    dtype = resolve_dtype(cfg.vlm.dtype, device)
+    sample = load_manifest(cfg.resolve(cfg.paths.manifest))[0]
+    history, _ = sample.load_trajectory()
+
+    processor, vlm = load_vlm(cfg.resolve(cfg.paths.model_dir), cfg.vlm, device, dtype)
+    vlm.eval()
+    baseline_text = generate_reasoning(
+        vlm,
+        processor,
+        sample.load_image(),
+        build_vlm_prefix(history),
+        device,
+        dtype,
+        cfg.vlm.max_reasoning_tokens,
+    )
+    prefix = extract_prefix_kv(
+        vlm,
+        processor,
+        sample.load_image(),
+        build_expert_prefix(history, sample.reasoning),
+        device,
+        dtype,
+    )
+
+    stage2 = cfg.stage2
+    architecture = ExpertArchitecture(
+        hidden_size=stage2.expert_hidden_size,
+        intermediate_size=stage2.expert_intermediate_size,
+        num_layers=stage2.expert_num_layers,
+        num_attention_heads=stage2.expert_num_attention_heads,
+        num_key_value_heads=stage2.expert_num_key_value_heads,
+        head_dim=stage2.expert_head_dim,
+        action_projection_hidden_size=stage2.action_projection_hidden_size,
+        action_fourier_features=stage2.action_fourier_features,
+        non_causal_attention=stage2.expert_non_causal_attention,
+    )
+    expert = KVCachedDiffusionExpert(vlm_text_config(vlm), architecture).to(
+        device=device, dtype=dtype
+    )
+    with torch.no_grad():
+        predicted_velocity = expert(
+            noisy_action=torch.randn(1, cfg.data.expert_waypoints, 2, device=device, dtype=dtype),
+            timestep=torch.full((1, 1, 1), 0.5, device=device, dtype=dtype),
+            prefix=prefix,
+        )
+
+    trainability = configure_stage1_trainability(vlm, cfg.stage1)
+    result = {
+        "device": str(device),
+        "dtype": str(dtype),
+        "vlm_parameters": parameter_counts(vlm),
+        "stage1_trainability": trainability,
+        "expert_parameters": parameter_counts(expert),
+        "expert_architecture": asdict(architecture),
+        "prefix_layers": prefix.num_layers,
+        "prefix_sequence_length": prefix.sequence_length,
+        "prefix_key_shape": list(prefix.keys[0].shape),
+        "expert_output_shape": list(predicted_velocity.shape),
+        "expert_output_finite": bool(torch.isfinite(predicted_velocity).all().item()),
+        "baseline_text": baseline_text,
+        "target_text": sample.reasoning,
+    }
+    write_json(cfg.root / "outputs" / "connection_test.json", result)
+    print(result)
+    del expert, vlm
+    empty_device_cache(device)
+
+
+if __name__ == "__main__":
+    main()