seeker
107 lines
3.4 KiB
Python
107 lines
3.4 KiB
Python
#!/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()
|