File size: 5,230 Bytes
2ea2835 | 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 141 142 143 | 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,
)
|