| from pathlib import Path |
| from types import SimpleNamespace |
|
|
| import pytest |
| import torch |
|
|
| from trainer.distillation import Trainer |
|
|
|
|
| def _branch(): |
| return { |
| "repulsion_denominator": torch.tensor(65536), |
| "fixed_teacher_attraction_denominator": torch.tensor(256), |
| "current_teacher_attraction_denominator": torch.tensor(32768), |
| "current_teacher_self_kernel_denominator": torch.tensor(16384), |
| "full_student_partial_gradient": torch.tensor(1), |
| } |
|
|
|
|
| def _trainer(): |
| trainer = object.__new__(Trainer) |
| trainer.fresh_video_total_samples = 256 |
| trainer.fresh_video_generator_grad_global_samples = 128 |
| trainer.fresh_video_full_student_partial_gradient = True |
| trainer.fresh_video_generator_grad_sample_scale = 2.0 |
| trainer.is_main_process = False |
| trainer.step = 0 |
| trainer.config = SimpleNamespace( |
| partial_mmd_normalization_identity=( |
| "full_student_population_denominators_v1"), |
| partial_mmd_kss_denominator=65536, |
| partial_mmd_fixed_attraction_denominator=256, |
| partial_mmd_live_cross_denominator=32768, |
| partial_mmd_teacher_self_denominator=16384, |
| partial_mmd_gradient_compensation=2.0, |
| ) |
| trainer.model = SimpleNamespace( |
| feature_type="triple_dinov3_vjepa2_videomae_v2_dense_latent", |
| current_teacher_attraction_enabled=True, |
| nystrom_repulsion_fifo_enabled=False, |
| nystrom_attraction_ema_enabled=False, |
| vjepa2_dense_layers=(15,), |
| dense_layer=15, |
| ) |
| return trainer |
|
|
|
|
| def _result(branch_names=("dinov3", "vjepa2", "videomae_v2")): |
| branches = {name: _branch() for name in branch_names} |
| metadata = { |
| "branches": branches, |
| "gradcache_kss_denominator": torch.tensor(65536), |
| "gradcache_full_student_partial_gradient": torch.tensor(1), |
| } |
| return SimpleNamespace( |
| metadata=metadata, |
| global_samples=256, |
| gradient_global_samples=128, |
| current_teacher_global_samples=128, |
| ) |
|
|
|
|
| def test_triple_geometry_validator_accepts_exact_three_branch_contract(): |
| _trainer()._validate_partial_mmd_geometry(_result()) |
|
|
|
|
| def test_triple_geometry_validator_rejects_legacy_dual_branch_set(): |
| with pytest.raises(RuntimeError, match="branch set mismatch"): |
| _trainer()._validate_partial_mmd_geometry( |
| _result(("vjepa2", "videomae_v2"))) |
|
|
|
|
| def test_release_config_matches_execution_geometry(): |
| import yaml |
| from utils.config import resolve_config_aliases |
| config = resolve_config_aliases(yaml.safe_load((Path(__file__).resolve().parents[2] / "configs/train.yaml").read_text())) |
| assert config["max_steps"] == 80 |
| assert config["lr_schedule_total_steps"] == 200 |
| assert config["triple_active_branches"] == ["dinov3", "vjepa2", "videomae_v2"] |
| assert config["triple_dinov3_weight"] == 10.0 |
| assert config["matched_manifest_samples"] == 8000 |
| assert config["triple_aggregation_divisor"] == 3.0 |
| assert config["batch_size"] == 2 |
| assert config["fresh_video_logical_micro_batch_size"] == 16 |
| assert config["fresh_video_physical_replay_batch_size"] == 1 |
| assert config["partial_mmd_geometry"] == [256, 128, 128] |
| assert config["optimizer_state_cpu_offload_during_backward"] |
| assert config["generator_replay_activation_cpu_offload"] |
|
|
|
|
| def test_logical_cotangent_is_sliced_without_rescaling_for_b1_replay(): |
| trainer = object.__new__(Trainer) |
| trainer.fresh_video_physical_replay_batch_size = 1 |
| trainer.fresh_video_logical_micro_batch_size = 2 |
| descriptor = { |
| "rng_state": {"opaque": True}, |
| "latent_shape": (2, 3, 1, 1, 1), |
| "generation_descriptor": { |
| "version": 2, |
| "num_generated_blocks": 1, |
| "noise": torch.arange(6).reshape(2, 3, 1, 1, 1), |
| "trajectory_seeds": torch.tensor([11, 22]), |
| }, |
| "student_latent": torch.arange(6).reshape(2, 3, 1, 1, 1), |
| } |
| physical_entry = { |
| "replay_descriptor": descriptor, |
| "image_or_video_shape": [2, 3, 1, 1, 1], |
| "conditional_dict": { |
| "prompt_embeds": torch.arange(8).reshape(2, 4), |
| }, |
| "unconditional_dict": { |
| "prompt_embeds": torch.arange(8, 16).reshape(2, 4), |
| }, |
| } |
| logical_entry = { |
| "features": {"vjepa2": torch.arange(12).reshape(2, 2, 3)}, |
| "physical_entries": (physical_entry,), |
| } |
| cotangent = { |
| "vjepa2": torch.arange(12, 24).reshape(2, 2, 3), |
| } |
| units = trainer._split_fresh_video_replay_units( |
| logical_entry, cotangent) |
| assert len(units) == 2 |
| for index, unit in enumerate(units): |
| assert unit["replay_descriptor"]["latent_shape"][0] == 1 |
| assert unit["image_or_video_shape"][0] == 1 |
| torch.testing.assert_close( |
| unit["cotangent"]["vjepa2"], |
| cotangent["vjepa2"][index:index + 1], |
| rtol=0, |
| atol=0, |
| ) |
| torch.testing.assert_close( |
| unit["replay_descriptor"]["generation_descriptor"][ |
| "trajectory_seeds"], |
| descriptor["generation_descriptor"][ |
| "trajectory_seeds"][index:index + 1], |
| rtol=0, |
| atol=0, |
| ) |
|
|
