| """Estimate Theorem-3 gradient second moments with disjoint stratified draws.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import random |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from omegaconf import OmegaConf |
| from scipy.stats import kendalltau, spearmanr |
|
|
| from starVLA.dataloader import build_dataloader |
| from starVLA.model.framework.base_framework import build_framework |
| from starVLA.model.framework.share_tools import apply_config_compat |
| from starVLA.scoring.flow_residual import ( |
| randomized_antithetic_stratified_flow_samples, |
| ) |
| from starVLA.scoring.partial_gradient import ( |
| activate_partial_gradient_group, |
| define_partial_gradient_groups, |
| finite_gradient_l2_norm, |
| ) |
| from starVLA.scoring.utility_audit import ( |
| alignment_audit, |
| leave_one_out_robustness, |
| ) |
|
|
|
|
| TRACE_IDENTITY_FIELDS = ( |
| "manifest_sha256", |
| "epoch", |
| "global_draw_id", |
| "dataset_index", |
| "dataset_name", |
| "trajectory_index", |
| "trajectory_id", |
| "anchor_index", |
| "base_index", |
| ) |
|
|
|
|
| def _sha256_file(path: Path) -> str: |
| digest = hashlib.sha256() |
| with path.open("rb") as source: |
| for chunk in iter(lambda: source.read(8 * 1024 * 1024), b""): |
| digest.update(chunk) |
| return digest.hexdigest() |
|
|
|
|
| def _set_random_seed(seed: int) -> None: |
| bounded_seed = int(seed) % (2**32) |
| random.seed(bounded_seed) |
| np.random.seed(bounded_seed) |
| torch.manual_seed(bounded_seed) |
| torch.cuda.manual_seed_all(bounded_seed) |
|
|
|
|
| def _trace_identity(trace: dict) -> dict: |
| return {field: trace.get(field) for field in TRACE_IDENTITY_FIELDS} |
|
|
|
|
| def _load_reference_artifact( |
| path: Path, |
| max_samples: int, |
| ) -> tuple[dict, list[dict]]: |
| payload = json.loads(path.read_text(encoding="utf-8")) |
| anchors = payload.get("anchors") |
| if not isinstance(anchors, list) or len(anchors) < max_samples: |
| raise ValueError("reference artifact does not contain enough anchors") |
| anchors = anchors[:max_samples] |
| return { |
| "path": str(path.resolve()), |
| "sha256": _sha256_file(path), |
| "checkpoint_sha256": payload.get("checkpoint_sha256"), |
| "old_gradient_seeds": payload.get("seed_separation", {}).get( |
| "gradient_seeds" |
| ), |
| }, anchors |
|
|
|
|
| def _load_aligned_examples(cfg, anchors: list[dict]) -> list[dict]: |
| dataloader = build_dataloader(cfg, dataset_py=cfg.datasets.vla_data.dataset_py) |
| examples = [] |
| for batch in dataloader: |
| examples.extend(batch) |
| if len(examples) >= len(anchors): |
| break |
| examples = examples[: len(anchors)] |
| if len(examples) != len(anchors): |
| raise RuntimeError("dataloader ended before audit anchors were recovered") |
| for index, (example, anchor) in enumerate(zip(examples, anchors)): |
| observed = _trace_identity(example.get("sampling_trace") or {}) |
| expected = _trace_identity(anchor["sampling_trace"]) |
| if observed != expected: |
| raise RuntimeError( |
| f"sample trace mismatch at position {index}: " |
| f"observed={observed}, expected={expected}" |
| ) |
| return examples |
|
|
|
|
| def _prepare_conditioning(model, example: dict): |
| images = [example["image"]] |
| instructions = [example["lang"]] |
| raw_actions = [example["action"]] |
| raw_states = [example["state"]] if "state" in example else None |
| with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): |
| vlm_inputs = model.qwen_vl_interface.build_qwenvl_inputs( |
| images=images, |
| instructions=instructions, |
| ) |
| vlm_outputs = model.qwen_vl_interface( |
| **vlm_inputs, |
| output_attentions=False, |
| output_hidden_states=True, |
| return_dict=True, |
| ) |
| hidden_states = vlm_outputs.hidden_states[-1].detach() |
| actions = torch.as_tensor( |
| np.asarray(raw_actions), |
| device=hidden_states.device, |
| dtype=hidden_states.dtype, |
| )[:, -model.action_horizon :, :] |
| states = ( |
| torch.as_tensor( |
| np.asarray(raw_states), |
| device=hidden_states.device, |
| dtype=hidden_states.dtype, |
| ) |
| if raw_states is not None |
| else None |
| ) |
| return hidden_states, actions, states |
|
|
|
|
| def _flow_loss( |
| model, |
| conditioning, |
| noise: torch.Tensor, |
| time_value: torch.Tensor, |
| ) -> torch.Tensor: |
| hidden_states, actions, states = conditioning |
| time = time_value.reshape(1) |
| expanded_time = time[:, None, None] |
| noisy_trajectory = (1.0 - expanded_time) * noise + expanded_time * actions |
| target_velocity = actions - noise |
| with torch.autocast("cuda", dtype=torch.float32): |
| prediction = model.action_model.predict_velocity( |
| hidden_states, |
| noisy_trajectory, |
| time, |
| state=states, |
| ) |
| return (prediction - target_velocity).square().mean() |
|
|
|
|
| def _measure_gradient_second_moment( |
| model, |
| examples: list[dict], |
| *, |
| estimator_seeds: list[int], |
| samples_per_estimator: int, |
| parameter_mode: str, |
| frozen_backbone_eval: bool, |
| ) -> tuple[np.ndarray, np.ndarray, dict]: |
| action_config = model.action_model.config |
| groups = define_partial_gradient_groups(model.action_model) |
| estimator_values = np.empty( |
| (len(examples), len(estimator_seeds)), |
| dtype=np.float64, |
| ) |
| draw_norms = np.empty( |
| (len(examples), len(estimator_seeds), samples_per_estimator), |
| dtype=np.float64, |
| ) |
| torch.cuda.empty_cache() |
| torch.cuda.reset_peak_memory_stats() |
| started = time.perf_counter() |
| parameters = [] |
|
|
| for anchor_index, example in enumerate(examples): |
| for estimator_index, estimator_seed in enumerate(estimator_seeds): |
| model.requires_grad_(False) |
| model.train() |
| if frozen_backbone_eval: |
| model.qwen_vl_interface.eval() |
| conditioning_seed = estimator_seed * 1_000_003 + anchor_index |
| _set_random_seed(conditioning_seed) |
| conditioning = _prepare_conditioning(model, example) |
| actions = conditioning[1] |
| noises, times = randomized_antithetic_stratified_flow_samples( |
| actions, |
| sample_count=samples_per_estimator, |
| seed=estimator_seed * 1_000_033 + anchor_index, |
| beta_alpha=float(action_config.noise_beta_alpha), |
| beta_beta=float(action_config.noise_beta_beta), |
| noise_s=float(action_config.noise_s), |
| ) |
| if parameter_mode == "full": |
| model.action_model.requires_grad_(True) |
| parameters = list(model.action_model.parameters()) |
| elif parameter_mode == "decoder": |
| parameters = activate_partial_gradient_group( |
| model.action_model, |
| groups["decoder"], |
| ) |
| else: |
| raise ValueError(f"unknown parameter mode: {parameter_mode}") |
|
|
| for draw_index in range(samples_per_estimator): |
| |
| |
| dropout_seed = ( |
| estimator_seed * 1_000_037 |
| + anchor_index * 101 |
| + draw_index // 2 |
| ) |
| _set_random_seed(dropout_seed) |
| model.zero_grad(set_to_none=True) |
| loss = _flow_loss( |
| model, |
| conditioning, |
| noises[draw_index], |
| times[draw_index], |
| ) |
| if loss.ndim != 0 or not torch.isfinite(loss): |
| raise FloatingPointError("flow loss must be a finite scalar") |
| loss.backward() |
| draw_norms[ |
| anchor_index, |
| estimator_index, |
| draw_index, |
| ] = finite_gradient_l2_norm(parameters) |
| estimator_values[anchor_index, estimator_index] = np.sqrt( |
| np.mean( |
| np.square( |
| draw_norms[anchor_index, estimator_index] |
| ) |
| ) |
| ) |
| print( |
| f"mode={parameter_mode} anchor={anchor_index + 1}/{len(examples)} " |
| f"h={np.sqrt(np.mean(np.square(draw_norms[anchor_index]))):.6f}", |
| flush=True, |
| ) |
| model.zero_grad(set_to_none=True) |
| torch.cuda.synchronize() |
| elapsed_seconds = time.perf_counter() - started |
| return estimator_values, draw_norms, { |
| "elapsed_seconds": elapsed_seconds, |
| "seconds_per_anchor": elapsed_seconds / len(examples), |
| "seconds_per_draw": ( |
| elapsed_seconds |
| / len(examples) |
| / len(estimator_seeds) |
| / samples_per_estimator |
| ), |
| "peak_allocated_bytes": torch.cuda.max_memory_allocated(), |
| "peak_reserved_bytes": torch.cuda.max_memory_reserved(), |
| "parameter_mode": parameter_mode, |
| "frozen_backbone_eval": frozen_backbone_eval, |
| "parameter_count": sum( |
| parameter.numel() for parameter in parameters |
| ), |
| } |
|
|
|
|
| def _repeatability(estimator_values: np.ndarray) -> dict: |
| if estimator_values.ndim != 2 or estimator_values.shape[1] < 4: |
| raise ValueError("at least four estimator columns are required") |
| if estimator_values.shape[1] % 2: |
| raise ValueError("estimator count must be even") |
| half = estimator_values.shape[1] // 2 |
| first = np.sqrt(np.mean(np.square(estimator_values[:, :half]), axis=1)) |
| second = np.sqrt(np.mean(np.square(estimator_values[:, half:]), axis=1)) |
| correlation = spearmanr(first, second) |
| concordance = kendalltau(first, second) |
| return { |
| "spearman": float(correlation.statistic), |
| "spearman_pvalue_asymptotic": float(correlation.pvalue), |
| "kendall": float(concordance.statistic), |
| "kendall_pvalue_asymptotic": float(concordance.pvalue), |
| "mean_anchor_cv": float( |
| np.mean( |
| estimator_values.std(axis=1, ddof=1) |
| / estimator_values.mean(axis=1) |
| ) |
| ), |
| } |
|
|
|
|
| def _gate( |
| sample_count: int, |
| proxy_repeatability: dict, |
| target_repeatability: dict, |
| alignment: dict, |
| robustness: dict, |
| ) -> dict: |
| checks = { |
| "sample_count_at_least_64": sample_count >= 64, |
| "target_repeatability_at_least_0_80": ( |
| target_repeatability["spearman"] >= 0.80 |
| ), |
| "proxy_repeatability_at_least_0_90": ( |
| proxy_repeatability["spearman"] >= 0.90 |
| ), |
| "proxy_target_spearman_at_least_0_70": ( |
| alignment["spearman_partial_gradient_vs_gradient_rms"] >= 0.70 |
| ), |
| "spearman_permutation_p_at_most_0_01": ( |
| alignment["random_permutation_control"][ |
| "one_sided_pvalue_rho_at_least_observed" |
| ] |
| <= 0.01 |
| ), |
| "top25_overlap_at_least_0_625": alignment["top25_overlap"] >= 0.625, |
| "second_moment_ratio_at_most_0_98": ( |
| alignment["proposal"]["information_to_baseline_ratio"] <= 0.98 |
| ), |
| "second_moment_permutation_p_at_most_0_01": ( |
| alignment["random_permutation_control"][ |
| "one_sided_pvalue_second_moment_at_most_observed" |
| ] |
| <= 0.01 |
| ), |
| "leave_one_out_spearman_all_positive": robustness[ |
| "all_spearman_positive" |
| ], |
| "leave_one_out_second_moment_all_below_one": robustness[ |
| "all_second_moment_ratios_below_one" |
| ], |
| } |
| return { |
| "checks": checks, |
| "pass": all(checks.values()), |
| "eligible_for_500_step_ablation": all(checks.values()), |
| } |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--config", type=Path, required=True) |
| parser.add_argument("--checkpoint", type=Path, required=True) |
| parser.add_argument("--reference-artifact", type=Path, required=True) |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument("--max-samples", type=int, default=8) |
| parser.add_argument("--samples-per-estimator", type=int, default=8) |
| parser.add_argument( |
| "--proxy-seeds", |
| type=int, |
| nargs="+", |
| default=[12300, 12301, 12302, 12303], |
| ) |
| parser.add_argument( |
| "--target-seeds", |
| type=int, |
| nargs="+", |
| default=[13300, 13301, 13302, 13303], |
| ) |
| parser.add_argument("--alpha", type=float, default=0.2) |
| parser.add_argument("--percentile-epsilon", type=float, default=0.1) |
| parser.add_argument("--permutations", type=int, default=10_000) |
| parser.add_argument("--permutation-seed", type=int, default=14100) |
| parser.add_argument("--frozen-backbone-eval", action="store_true") |
| args = parser.parse_args() |
| if args.max_samples < 4: |
| raise ValueError("max-samples must be at least four") |
| for name, seeds in ( |
| ("proxy-seeds", args.proxy_seeds), |
| ("target-seeds", args.target_seeds), |
| ): |
| if len(seeds) < 4 or len(seeds) % 2 or len(set(seeds)) != len(seeds): |
| raise ValueError(f"{name} must contain an even count of unique seeds") |
| if set(args.proxy_seeds) & set(args.target_seeds): |
| raise ValueError("proxy and target seeds must be disjoint") |
| if args.samples_per_estimator < 2 or args.samples_per_estimator % 2: |
| raise ValueError("samples-per-estimator must be a positive even integer") |
|
|
| reference, reference_anchors = _load_reference_artifact( |
| args.reference_artifact, |
| args.max_samples, |
| ) |
| checkpoint_sha256 = _sha256_file(args.checkpoint) |
| if ( |
| reference["checkpoint_sha256"] is not None |
| and reference["checkpoint_sha256"] != checkpoint_sha256 |
| ): |
| raise RuntimeError("checkpoint hash differs from reference artifact") |
| all_old_seeds = set(reference["old_gradient_seeds"] or []) |
| if all_old_seeds & (set(args.proxy_seeds) | set(args.target_seeds)): |
| raise ValueError("new seeds overlap the old gradient audit") |
|
|
| cfg = apply_config_compat(OmegaConf.load(args.config)) |
| cfg.output_dir = str(args.output.parent / "stratified_gradient_runtime") |
| examples = _load_aligned_examples(cfg, reference_anchors) |
| model = build_framework(cfg) |
| checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=True) |
| model.load_state_dict(checkpoint, strict=True) |
| model.to("cuda") |
|
|
| proxy_estimators, proxy_draws, proxy_runtime = ( |
| _measure_gradient_second_moment( |
| model, |
| examples, |
| estimator_seeds=args.proxy_seeds, |
| samples_per_estimator=args.samples_per_estimator, |
| parameter_mode="decoder", |
| frozen_backbone_eval=args.frozen_backbone_eval, |
| ) |
| ) |
| target_estimators, target_draws, target_runtime = ( |
| _measure_gradient_second_moment( |
| model, |
| examples, |
| estimator_seeds=args.target_seeds, |
| samples_per_estimator=args.samples_per_estimator, |
| parameter_mode="full", |
| frozen_backbone_eval=args.frozen_backbone_eval, |
| ) |
| ) |
| proxy_scores = np.sqrt(np.mean(np.square(proxy_estimators), axis=1)) |
| target_gradient_rms = np.sqrt( |
| np.mean(np.square(target_estimators), axis=1) |
| ) |
| old_gradient_rms = np.asarray( |
| [anchor["gradient_rms"] for anchor in reference_anchors], |
| dtype=np.float64, |
| ) |
| mean_losses = np.asarray( |
| [anchor["mean_gradient_loss"] for anchor in reference_anchors], |
| dtype=np.float64, |
| ) |
| proxy_repeatability = _repeatability(proxy_estimators) |
| target_repeatability = _repeatability(target_estimators) |
| alignment = alignment_audit( |
| proxy_scores, |
| target_gradient_rms, |
| mean_losses, |
| alpha=args.alpha, |
| percentile_epsilon=args.percentile_epsilon, |
| permutation_count=args.permutations, |
| permutation_seed=args.permutation_seed, |
| score_label="partial_gradient", |
| ) |
| robustness = leave_one_out_robustness( |
| proxy_scores, |
| target_gradient_rms, |
| alpha=args.alpha, |
| percentile_epsilon=args.percentile_epsilon, |
| ) |
| gate = _gate( |
| len(examples), |
| proxy_repeatability, |
| target_repeatability, |
| alignment, |
| robustness, |
| ) |
| old_target_correlation = spearmanr( |
| target_gradient_rms, |
| old_gradient_rms, |
| ) |
| payload = { |
| "schema_version": 1, |
| "status": ( |
| "formal_64_anchor_stratified_gradient_gate" |
| if len(examples) >= 64 |
| else "stratified_gradient_smoke_not_a_gate_certificate" |
| ), |
| "config": str(args.config.resolve()), |
| "checkpoint": str(args.checkpoint.resolve()), |
| "checkpoint_sha256": checkpoint_sha256, |
| "reference_artifact": reference, |
| "sample_count": len(examples), |
| "estimator": { |
| "definition": "sqrt(mean over xi of squared per-draw gradient L2 norm)", |
| "randomized_beta_stratification": True, |
| "antithetic_gaussian_noise": True, |
| "common_dropout_within_antithetic_pair": True, |
| "samples_per_estimator": args.samples_per_estimator, |
| "proxy_seeds": args.proxy_seeds, |
| "target_seeds": args.target_seeds, |
| "seeds_disjoint": True, |
| "frozen_backbone_eval": args.frozen_backbone_eval, |
| }, |
| "proxy_runtime": proxy_runtime, |
| "target_runtime": target_runtime, |
| "proxy_repeatability": proxy_repeatability, |
| "target_repeatability": target_repeatability, |
| "new_vs_old_target_spearman": float(old_target_correlation.statistic), |
| "new_vs_old_target_spearman_pvalue": float(old_target_correlation.pvalue), |
| "alignment": alignment, |
| "leave_one_out_robustness": robustness, |
| "gate": gate, |
| "eligible_for_500_step_ablation": gate[ |
| "eligible_for_500_step_ablation" |
| ], |
| "anchors": [ |
| { |
| "anchor_position": index, |
| "sampling_trace": reference_anchors[index]["sampling_trace"], |
| "proxy_estimators": proxy_estimators[index].tolist(), |
| "proxy_draw_norms": proxy_draws[index].tolist(), |
| "proxy_gradient_rms": float(proxy_scores[index]), |
| "target_estimators": target_estimators[index].tolist(), |
| "target_draw_norms": target_draws[index].tolist(), |
| "target_gradient_rms": float(target_gradient_rms[index]), |
| "old_four_draw_gradient_rms": float(old_gradient_rms[index]), |
| } |
| for index in range(len(examples)) |
| ], |
| } |
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| temporary = args.output.with_suffix(args.output.suffix + ".tmp") |
| temporary.write_text(json.dumps(payload, indent=2), encoding="utf-8") |
| temporary.replace(args.output) |
| print(json.dumps(payload, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
|
|
|
|
| def _measure_gradient_second_moment_draw_shard( |
| model, |
| examples: list[dict], |
| *, |
| estimator_seeds: list[int], |
| samples_per_estimator: int, |
| draw_indices: list[int], |
| parameter_mode: str, |
| frozen_backbone_eval: bool, |
| ) -> tuple[np.ndarray, np.ndarray, dict]: |
| action_config = model.action_model.config |
| groups = define_partial_gradient_groups(model.action_model) |
| values = np.empty((len(examples), len(estimator_seeds)), dtype=np.float64) |
| draw_norms = np.empty( |
| (len(examples), len(estimator_seeds), len(draw_indices)), |
| dtype=np.float64, |
| ) |
| torch.cuda.empty_cache() |
| torch.cuda.reset_peak_memory_stats() |
| started = time.perf_counter() |
| parameters = [] |
| for anchor_index, example in enumerate(examples): |
| for estimator_index, estimator_seed in enumerate(estimator_seeds): |
| model.requires_grad_(False) |
| model.train() |
| if frozen_backbone_eval: |
| model.qwen_vl_interface.eval() |
| conditioning_seed = estimator_seed * 1_000_003 + anchor_index |
| _set_random_seed(conditioning_seed) |
| conditioning = _prepare_conditioning(model, example) |
| actions = conditioning[1] |
| noises, times = randomized_antithetic_stratified_flow_samples( |
| actions, |
| sample_count=samples_per_estimator, |
| seed=estimator_seed * 1_000_033 + anchor_index, |
| beta_alpha=float(action_config.noise_beta_alpha), |
| beta_beta=float(action_config.noise_beta_beta), |
| noise_s=float(action_config.noise_s), |
| ) |
| if parameter_mode != "full": |
| raise ValueError(f"unknown parameter mode: {parameter_mode}") |
| model.action_model.requires_grad_(True) |
| parameters = list(model.action_model.parameters()) |
| for output_index, draw_index in enumerate(draw_indices): |
| dropout_seed = ( |
| estimator_seed * 1_000_037 |
| + anchor_index * 101 |
| + draw_index // 2 |
| ) |
| _set_random_seed(dropout_seed) |
| model.zero_grad(set_to_none=True) |
| loss = _flow_loss( |
| model, |
| conditioning, |
| noises[draw_index], |
| times[draw_index], |
| ) |
| if loss.ndim != 0 or not torch.isfinite(loss): |
| raise FloatingPointError("flow loss must be a finite scalar") |
| loss.backward() |
| draw_norms[anchor_index, estimator_index, output_index] = finite_gradient_l2_norm(parameters) |
| values[anchor_index, estimator_index] = np.sqrt( |
| np.mean(np.square(draw_norms[anchor_index, estimator_index])) |
| ) |
| print( |
| f"mode={parameter_mode} anchor={anchor_index + 1}/{len(examples)} " |
| f"h={np.sqrt(np.mean(np.square(draw_norms[anchor_index]))):.6f}", |
| flush=True, |
| ) |
| model.zero_grad(set_to_none=True) |
| torch.cuda.synchronize() |
| elapsed_seconds = time.perf_counter() - started |
| return values, draw_norms, { |
| "elapsed_seconds": elapsed_seconds, |
| "seconds_per_anchor": elapsed_seconds / len(examples), |
| "seconds_per_draw": elapsed_seconds / len(examples) / len(estimator_seeds) / len(draw_indices), |
| "peak_allocated_bytes": torch.cuda.max_memory_allocated(), |
| "peak_reserved_bytes": torch.cuda.max_memory_reserved(), |
| "parameter_mode": parameter_mode, |
| "frozen_backbone_eval": frozen_backbone_eval, |
| "parameter_count": sum(parameter.numel() for parameter in parameters), |
| } |
|
|
