diff --git a/code/flash-linear-attention/tests/models/test_modeling_rwkv7.py b/code/flash-linear-attention/tests/models/test_modeling_rwkv7.py new file mode 100644 index 0000000000000000000000000000000000000000..5f377164fae141c2633bbcc7f8089fa0ce0472bc --- /dev/null +++ b/code/flash-linear-attention/tests/models/test_modeling_rwkv7.py @@ -0,0 +1,57 @@ + +import pytest +import torch + +from fla.models import RWKV7Config + +from .test_modeling_base import run_test_generation, run_test_model_forward_backward + + +# =================================================================================== +# Test for Modeling (Forward/Backward Pass) +# =================================================================================== +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ], +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, RWKV7Config, use_l2warp=use_l2warp, dtype=dtype) + +# =================================================================================== +# Test for Generation +# =================================================================================== + + +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ], +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): + run_test_generation(L, B, T, H, D, RWKV7Config, dtype) diff --git a/code/flash-linear-attention/tests/models/test_modeling_samba.py b/code/flash-linear-attention/tests/models/test_modeling_samba.py new file mode 100644 index 0000000000000000000000000000000000000000..77a236c0c32b508736734fc5770b7fb211063576 --- /dev/null +++ b/code/flash-linear-attention/tests/models/test_modeling_samba.py @@ -0,0 +1,56 @@ + +import pytest +import torch + +from fla.models import SambaConfig + +from .test_modeling_base import run_test_generation, run_test_model_forward_backward + + +# =================================================================================== +# Test for Modeling (Forward/Backward Pass) +# =================================================================================== +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 18, 64, True, torch.bfloat16), + (4, 4, 1024, 18, 64, False, torch.bfloat16), + (4, 4, 1024, 9, 128, False, torch.bfloat16), + ] + ], +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, SambaConfig, use_l2warp=use_l2warp, dtype=dtype) + + +# =================================================================================== +# Test for Generation +# =================================================================================== +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 18, 64, torch.float16), + ] + ], +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): + run_test_generation(L, B, T, H, D, SambaConfig, dtype) diff --git a/code/flash-linear-attention/tests/models/test_modeling_transformer.py b/code/flash-linear-attention/tests/models/test_modeling_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..b704a72a1e0b08f72d812eef7f2401b7dcc34d56 --- /dev/null +++ b/code/flash-linear-attention/tests/models/test_modeling_transformer.py @@ -0,0 +1,56 @@ + +import pytest +import torch + +from fla.models import TransformerConfig + +from .test_modeling_base import run_test_generation, run_test_model_forward_backward + + +# =================================================================================== +# Test for Modeling (Forward/Backward Pass) +# =================================================================================== +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ], +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, TransformerConfig, use_l2warp=use_l2warp, dtype=dtype) + + +# =================================================================================== +# Test for Generation +# =================================================================================== +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ], +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): + run_test_generation(L, B, T, H, D, TransformerConfig, dtype) diff --git a/code/flash-linear-attention/tests/models/test_modeling_utils.py b/code/flash-linear-attention/tests/models/test_modeling_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e5df066fe470c2b761b6da40b221e11a185a8806 --- /dev/null +++ b/code/flash-linear-attention/tests/models/test_modeling_utils.py @@ -0,0 +1,84 @@ + +import math + +import torch +import torch.nn as nn +from transformers import AutoModelForCausalLM + +from fla.utils import device + +# Models that do not yet support variable sequence lengths (for modeling tests) +MODELING_UNSUPPORTED_VARLEN = [ + "ABCConfig", "ForgettingTransformerConfig", "LinearAttentionConfig", "LightNetConfig", + "Mamba2Config", "MambaConfig", "MesaNetConfig", "SambaConfig", + "RodimusConfig", +] + +# Models not yet ready for basic testing +NOT_READY_FOR_TESTING = ['RodimusConfig'] + +# Models requiring specific hardware (e.g., NVIDIA Hopper) +HOPPER_EXCLUSIVE = [] + +GENERATION_UNSUPPORTED = [ + "ABCConfig", "LinearAttentionConfig", "LightNetConfig", + "Mamba2Config", "MambaConfig", "NSAConfig", "SambaConfig", "RWKV6Config", "RWKV7Config", + "DeltaFormerConfig", +] + + +def create_model_and_config(config_class, L, H, D, dtype, **kwargs): + """ + A helper function to create a model and its configuration. + """ + config_params = { + 'hidden_size': H * D, + 'num_hidden_layers': L, + **({'num_heads': H} if config_class.__name__ != 'NSAConfig' else {}), + **kwargs, + } + config = config_class(**config_params) + model = AutoModelForCausalLM.from_config(config) + model.apply(init_weights_recursively) + model.to(dtype).to(device) + return model, config + + +def init_weights_with_asymmetric_pattern(module): + """Initialize weights with asymmetric patterns for debugging. + + Args: + module: The module to initialize weights for. + """ + if isinstance(module, (nn.Linear, nn.Conv1d)): + nn.init.kaiming_normal_(module.weight, a=math.sqrt(5)) + with torch.no_grad(): + shape = module.weight.shape + if len(shape) > 1: + quarter_size = shape[0] // 4 + module.weight[:quarter_size] *= 1.2 + module.weight[-quarter_size:] *= 0.8 + if shape[0] == shape[1]: + idx = torch.arange(min(shape[0], shape[1])) + module.weight[idx, idx] += 0.05 + if module.bias is not None: + fan_in, _ = nn.init._calculate_fan_in_and_fan_out(module.weight) + bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 + nn.init.uniform_(module.bias, -bound, bound) + with torch.no_grad(): + module.bias[::3] *= 1.1 + module.bias[1::3] *= 0.9 + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=0.02) + with torch.no_grad(): + vocab_size, dim = module.weight.shape + pattern = 0.01 * torch.sin(torch.arange(dim) * (6.28 / dim)) + for i in range(min(100, vocab_size)): + module.weight[i] += pattern * (1 + i % 5) * 0.2 + + +def init_weights_recursively(module): + if hasattr(module, 'weight'): + init_weights_with_asymmetric_pattern(module) + for submodule in module.children(): + init_weights_recursively(submodule) diff --git a/code/flash-linear-attention/tests/modules/test_activation.py b/code/flash-linear-attention/tests/modules/test_activation.py new file mode 100644 index 0000000000000000000000000000000000000000..2b980a4da2f93d1d4551afbc3c8f576e872b1ca9 --- /dev/null +++ b/code/flash-linear-attention/tests/modules/test_activation.py @@ -0,0 +1,133 @@ + +import pytest +import torch +import torch.nn.functional as F + +from fla.modules.activations import logsigmoid, sigmoid, swiglu, swiglu_linear, swish +from fla.utils import assert_close, device + + +@pytest.mark.parametrize( + ('B', 'T', 'D', 'compile'), + [ + (1, 1, 64, False), + (2, 500, 128, False), + (2, 512, 128, True), + (3, 2048, 1200, True), + ], +) +def test_sigmoid(B: int, T: int, D: int, compile: bool): + torch.manual_seed(42) + x = torch.randn(B, T, D, device=device, requires_grad=True) + y_ref = torch.sigmoid(x) + y_tri = sigmoid(x) if not compile else torch.compile(sigmoid)(x) + + g = torch.randn_like(y_ref) + dx_ref = torch.autograd.grad(y_ref, x, g)[0] + dx_tri = torch.autograd.grad(y_tri, x, g)[0] + + assert_close('sigmoid fwd', y_ref, y_tri, 1e-3) + assert_close('sigmoid bwd', dx_ref, dx_tri, 1e-3) + + +@pytest.mark.parametrize( + ('B', 'T', 'D', 'temperature', 'compile'), + [ + (1, 1, 64, 1.0, False), + (2, 500, 128, 0.5, False), + (2, 512, 128, 0.5, True), + (3, 2048, 1200, 2.0, True), + ], +) +def test_logsigmoid(B: int, T: int, D: int, temperature: float, compile: bool): + torch.manual_seed(42) + x = torch.randn(B, T, D, device=device, requires_grad=True) + y_ref = F.logsigmoid(x) / temperature + y_tri = logsigmoid(x, temperature) if not compile else torch.compile(logsigmoid)(x, temperature) + + g = torch.randn_like(y_ref) + dx_ref = torch.autograd.grad(y_ref, x, g)[0] + dx_tri = torch.autograd.grad(y_tri, x, g)[0] + + assert_close('logsigmoid fwd', y_ref, y_tri, 1e-3) + assert_close('logsigmoid bwd', dx_ref, dx_tri, 1e-3) + + +@pytest.mark.parametrize( + ('B', 'T', 'D', 'compile'), + [ + (1, 1, 64, True), + (2, 500, 128, True), + (2, 512, 128, False), + (3, 2048, 1200, False), + ], +) +def test_swish(B: int, T: int, D: int, compile: bool): + torch.manual_seed(42) + x = torch.randn(B, T, D, device=device, requires_grad=True) + y_ref = F.silu(x) + y_tri = swish(x) if not compile else torch.compile(swish)(x) + + g = torch.randn_like(y_ref) + dx_ref = torch.autograd.grad(y_ref, x, g)[0] + dx_tri = torch.autograd.grad(y_tri, x, g)[0] + + assert_close('swish fwd', y_ref, y_tri, 1e-3) + assert_close('swish bwd', dx_ref, dx_tri, 1e-3) + + +@pytest.mark.parametrize( + ('B', 'T', 'D', 'compile'), + [ + (1, 1, 64, True), + (2, 500, 128, True), + (2, 512, 128, False), + (3, 2048, 1200, False), + ], +) +def test_swiglu(B: int, T: int, D: int, compile: bool): + torch.manual_seed(42) + x = torch.randn(B, T, D, device=device, requires_grad=True) + y = torch.randn(B, T, D, device=device, requires_grad=True) + + y_ref = F.silu(x) * y + y_tri = swiglu(x, y) if not compile else torch.compile(swiglu)(x, y) + + g = torch.randn_like(y_ref) + dx_ref, dy_ref = torch.autograd.grad(y_ref, (x, y), g) + dx_tri, dy_tri = torch.autograd.grad(y_tri, (x, y), g) + + assert_close('swiglu fwd', y_ref, y_tri, 1e-3) + assert_close('swiglu dx', dx_ref, dx_tri, 1e-3) + assert_close('swiglu dy', dy_ref, dy_tri, 1e-3) + + +@pytest.mark.parametrize( + ('B', 'T', 'D', 'O', 'compile'), + [ + (2, 512, 128, 256, True), + (1, 1, 64, 32, False), + (2, 500, 128, 64, True), + (3, 2048, 1200, 600, False), + ], +) +def test_swiglu_linear(B: int, T: int, D: int, O: int, compile: bool): # noqa: E741 + torch.manual_seed(42) + x = torch.randn(B, T, D, device=device, requires_grad=True) + y = torch.randn(B, T, D, device=device, requires_grad=True) + w = torch.randn(O, D, device=device, requires_grad=True) + b = torch.randn(O, device=device, requires_grad=True) + + z_ref = F.silu(x) * y + out_ref = F.linear(z_ref, w, b) + out_tri = swiglu_linear(x, y, w, b) if not compile else torch.compile(swiglu_linear)(x, y, w, b) + + g = torch.randn_like(out_ref) + dx_ref, dy_ref, dw_ref, db_ref = torch.autograd.grad(out_ref, (x, y, w, b), g) + dx_tri, dy_tri, dw_tri, db_tri = torch.autograd.grad(out_tri, (x, y, w, b), g) + + assert_close('swiglu_linear out', out_ref, out_tri, 1e-3) + assert_close('swiglu_linear dx', dx_ref, dx_tri, 1e-3) + assert_close('swiglu_linear dy', dy_ref, dy_tri, 1e-3) + assert_close('swiglu_linear dw', dw_ref, dw_tri, 1e-3) + assert_close('swiglu_linear db', db_ref, db_tri, 1e-3) diff --git a/code/flash-linear-attention/tests/modules/test_conv.py b/code/flash-linear-attention/tests/modules/test_conv.py new file mode 100644 index 0000000000000000000000000000000000000000..dcd780018826ccc54e07e2e1e2551b42d6c686af --- /dev/null +++ b/code/flash-linear-attention/tests/modules/test_conv.py @@ -0,0 +1,713 @@ + +import pytest +import torch +import torch.nn.functional as F +from einops import rearrange + +from fla.modules.convolution import ShortConvolution, causal_conv1d, causal_conv1d_update +from fla.utils import assert_close, device + +try: + from causal_conv1d import causal_conv1d_fn +except ImportError: + causal_conv1d_fn = None + + +def causal_conv1d_ref_torch( + x, + weight, + bias=None, + initial_state=None, + output_final_state=False, + final_states_out=None, + activation=None, +): + """ + x: (batch, dim, seqlen) + weight: (dim, width) + bias: (dim,) + initial_state: (batch, dim, width - 1) + final_states_out: (batch, dim, width - 1) + + out: (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + dtype_in = x.dtype + x = x.to(weight.dtype) + seqlen = x.shape[-1] + dim, width = weight.shape + if initial_state is None: + out = F.conv1d(x, weight.unsqueeze(1), bias, padding=width - 1, groups=dim) + else: + x = torch.cat([initial_state, x], dim=-1) + out = F.conv1d(x, weight.unsqueeze(1), bias, padding=0, groups=dim) + out = out[..., :seqlen] + if output_final_state: + final_states = F.pad(x, (width - 1 - x.shape[-1], 0)).to( + dtype_in, + ) # (batch, dim, width - 1) + if final_states_out is not None: + final_states_out.copy_(final_states) + else: + final_states_out = final_states + out = (out if activation is None else F.silu(out)).to(dtype=dtype_in) + return out if not output_final_state else (out, final_states_out) + + +def causal_conv1d_update_ref_torch(x, conv_state, weight, bias=None, activation=None, cache_seqlens=None): + """ + x: (batch, dim) or (batch, dim, seqlen) + conv_state: (batch, dim, state_len), where state_len >= width - 1 + weight: (dim, width) + bias: (dim,) + cache_seqlens: (batch,), dtype int32. + If not None, the conv_state is treated as a circular buffer. + The conv_state will be updated by copying x to the conv_state starting at the index + @cache_seqlens % state_len before performing the convolution. + + out: (batch, dim) or (batch, dim, seqlen) + """ + if activation not in [None, "silu", "swish"]: + raise NotImplementedError("activation must be None, silu, or swish") + dtype_in = x.dtype + unsqueeze = x.dim() == 2 + if unsqueeze: + x = x.unsqueeze(-1) + batch, dim, seqlen = x.shape + width = weight.shape[1] + state_len = conv_state.shape[-1] + assert conv_state.shape == (batch, dim, state_len) + assert weight.shape == (dim, width) + if cache_seqlens is None: + x_new = torch.cat([conv_state, x], dim=-1).to(weight.dtype) # (batch, dim, state_len + seqlen) + conv_state.copy_(x_new[:, :, -state_len:]) + else: + width_idx = torch.arange(-(width - 1), 0, dtype=torch.long, device=x.device).unsqueeze(0) + cache_seqlens.unsqueeze(1) + width_idx = torch.remainder(width_idx, state_len).unsqueeze(1).expand(-1, dim, -1) + x_new = torch.cat([conv_state.gather(2, width_idx), x], dim=-1).to(weight.dtype) + copy_idx = torch.arange(seqlen, dtype=torch.long, device=x.device).unsqueeze(0) + cache_seqlens.unsqueeze(1) + copy_idx = torch.remainder(copy_idx, state_len).unsqueeze(1).expand(-1, dim, -1) + conv_state.scatter_(2, copy_idx, x) + out = F.conv1d(x_new, weight.unsqueeze(1), bias, padding=0, groups=dim)[:, :, -seqlen:] + if unsqueeze: + out = out.squeeze(-1) + return (out if activation is None else F.silu(out)).to(dtype=dtype_in) + + +@pytest.mark.parametrize( + ('B', 'T', 'D', 'W', 'activation', 'has_bias', 'has_residual', 'dtype'), + [ + pytest.param(*test, id="B{0}_T{1}_D{2}_W{3}_activation{4}_has_bias{5}_has_residual{6}_{7}".format(*test)) + for test in [ + (2, 64, 128, 3, "swish", True, True, torch.float32), + (2, 128, 128, 4, "swish", False, True, torch.float32), + (2, 64, 128, 3, "swish", True, False, torch.float32), + (2, 128, 128, 4, "swish", False, False, torch.float32), + (2, 500, 1024, 3, None, True, True, torch.float32), + (2, 1024, 1024, 4, None, False, True, torch.float32), + (2, 64, 128, 3, None, True, False, torch.float16), + (2, 128, 128, 4, None, False, False, torch.float16), + ] + ], +) +def test_conv( + B: int, + T: int, + D: int, + W: int, + activation: str, + has_bias: bool, + has_residual: bool, + dtype: torch.dtype, +): + torch.manual_seed(42) + + x = torch.randn(B, T, D).to(device, dtype).requires_grad_(True) + weight = torch.randn(D, W).to(device, dtype).requires_grad_(True) + bias = torch.randn(D).to(device, dtype).requires_grad_(True) if has_bias else None + residual = x.detach().clone().requires_grad_(True) if has_residual else None + dy = torch.randn(B, T, D).to(device, dtype) + + ref = causal_conv1d_ref_torch( + x=rearrange(x, "b t d -> b d t"), + weight=weight, + bias=bias, + activation=activation, + ) + ref = rearrange(ref, "b d t -> b t d") + if has_residual: + ref += residual + ref.backward(dy) + ref_dx, x.grad = x.grad, None + ref_dw, weight.grad = weight.grad, None + if has_bias: + ref_db, bias.grad = bias.grad, None + if has_residual: + ref_dr, residual.grad = residual.grad, None + + tri, _ = causal_conv1d(x, weight, bias, residual=residual, activation=activation) + tri.backward(dy) + tri_dx, x.grad = x.grad, None + tri_dw, weight.grad = weight.grad, None + if has_bias: + tri_db, bias.grad = bias.grad, None + if has_residual: + tri_dr, residual.grad = residual.grad, None + + assert_close(" y", ref, tri, 1e-3) + assert_close("dx", ref_dx, tri_dx, 1e-3) + assert_close("dw", ref_dw, tri_dw, 1e-3) + if has_bias: + assert_close("db", ref_db, tri_db, 1e-3) + if has_residual: + assert_close("dr", ref_dr, tri_dr, 1e-3) + + +@pytest.mark.parametrize( + ('N', 'T', 'D', 'W', 'activation', 'has_bias', 'has_residual', 'dtype'), + [ + pytest.param(*test, id="N{0}_T{1}_D{2}_W{3}_activation{4}_has_bias{5}_has_residual{6}_{7}".format(*test)) + for test in [ + (4, 500, 128, 3, "swish", True, True, torch.float32), + (4, 1024, 200, 4, "swish", False, True, torch.float32), + (4, 500, 128, 3, None, True, False, torch.float16), + (4, 1024, 1024, 4, None, False, False, torch.float16), + ] + ], +) +def test_conv_varlen( + N: int, + T: int, + D: int, + W: int, + activation: str, + has_bias: bool, + has_residual: bool, + dtype: torch.dtype, +): + torch.manual_seed(42) + cu_seqlens = torch.cat([ + torch.tensor([0], dtype=torch.long), + torch.arange(16, T)[torch.randperm(T - 16)[:N-1]], + torch.tensor([T], dtype=torch.long), + ], 0).to(device).sort()[0] + + x = torch.randn(1, T, D).to(device, dtype).requires_grad_(True) + weight = torch.randn(D, W).to(device, dtype).requires_grad_(True) + bias = torch.randn(D).to(device, dtype).requires_grad_(True) if has_bias else None + residual = x.detach().clone().requires_grad_(True) if has_residual else None + dy = torch.randn(1, T, D).to(device, dtype) + + ref = torch.cat([ + rearrange( + causal_conv1d_ref_torch( + x=rearrange(x[:, bos:eos].contiguous(), "b t d -> b d t"), + weight=weight, + bias=bias, + activation=activation, + ), + "b t d -> b d t", + ) + (residual[:, bos:eos] if has_residual else torch.zeros_like(x[:, bos:eos])) + for bos, eos in zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False) + ], 1) + ref.backward(dy) + ref_dx, x.grad = x.grad, None + ref_dw, weight.grad = weight.grad, None + if has_bias: + ref_db, bias.grad = bias.grad, None + if has_residual: + ref_dr, residual.grad = residual.grad, None + + tri, _ = causal_conv1d(x, weight, bias, residual=residual, activation=activation, cu_seqlens=cu_seqlens) + tri.backward(dy) + tri_dx, x.grad = x.grad, None + tri_dw, weight.grad = weight.grad, None + if has_bias: + tri_db, bias.grad = bias.grad, None + if has_residual: + tri_dr, residual.grad = residual.grad, None + + assert_close(" y", ref, tri, 1e-3) + assert_close("dx", ref_dx, tri_dx, 1e-3) + assert_close("dw", ref_dw, tri_dw, 1e-3) + if has_bias: + assert_close("db", ref_db, tri_db, 1e-3) + if has_residual: + assert_close("dr", ref_dr, tri_dr, 1e-3) + + +@pytest.mark.parametrize( + ('B', 'T', 'D', 'W', 'activation', 'has_bias', 'has_residual', 'dtype'), + [ + pytest.param(*test, id="B{0}_T{1}_D{2}_W{3}_activation{4}_has_bias{5}_has_residual{6}_{7}".format(*test)) + for test in [ + (2, 64, 128, 3, "swish", True, True, torch.float32), + (2, 128, 128, 4, "swish", False, True, torch.float32), + (2, 64, 128, 3, "swish", True, False, torch.float32), + (2, 128, 128, 4, "swish", False, False, torch.float32), + (2, 500, 1024, 3, None, True, True, torch.float32), + (2, 1024, 1024, 4, None, False, True, torch.float32), + (2, 64, 128, 3, None, True, False, torch.float16), + (2, 128, 128, 4, None, False, False, torch.float16), + ] + ], +) +@torch.no_grad +def test_conv_decoding( + B: int, + T: int, + D: int, + W: int, + activation: str, + has_bias: bool, + has_residual: bool, + dtype: torch.dtype, +): + torch.manual_seed(42) + + x = torch.randn(B, T, D).to(device, dtype) + weight = torch.randn(D, W).to(device, dtype) * 0 + bias = torch.randn(D).to(device, dtype) if has_bias else None + residual = x.clone() if has_residual else None + + ref = causal_conv1d_ref_torch( + x=rearrange(x, "b t d -> b d t"), + weight=weight, + bias=bias, + activation=activation, + ) + ref = rearrange(ref, "b d t -> b t d") + if has_residual: + ref += residual + ref_cache = x.new_zeros(B, D, W) + ref_cache[:, :, -min(W, T):].copy_(rearrange(x[..., -min(W, T):, :], 'n w d -> n d w')) + + tri = torch.zeros_like(x) + tri_cache = x.new_zeros(B, D, W) + for i in range(T): + y, tri_cache = causal_conv1d_update( + x=x[:, i:i+1, :], + cache=tri_cache, + residual=residual[:, i:i+1, :] if has_residual else None, + weight=weight, + bias=bias, + activation=activation, + ) + tri[:, i:i+1, :] = y + + assert_close(" y", ref, tri, 1e-3) + assert_close("cache", ref_cache, tri_cache, 1e-3) + + +@pytest.mark.parametrize( + ('B', 'T', 'D', 'W', 'activation', 'has_bias', 'has_residual', 'dtype', 'backend'), + [ + pytest.param( + *test, id="B{0}_T{1}_D{2}_W{3}_activation{4}_has_bias{5}_has_residual{6}_{7}_{8}".format(*test)) + for test in [ + (2, 64, 128, 3, "swish", True, True, torch.float32, 'triton'), + (2, 128, 128, 4, "swish", False, True, torch.float32, 'triton'), + (2, 64, 128, 3, "swish", True, False, torch.float32, 'triton'), + (2, 128, 128, 4, "swish", False, False, torch.float32, 'triton'), + (2, 500, 1024, 3, None, True, True, torch.float32, 'triton'), + (2, 1024, 1024, 4, None, False, True, torch.float32, 'triton'), + (2, 64, 128, 3, None, True, False, torch.float16, 'triton'), + (2, 128, 128, 4, None, False, False, torch.float16, 'triton'), + (2, 64, 128, 3, "swish", True, True, torch.float32, 'cuda'), + (2, 128, 128, 4, "swish", False, True, torch.float32, 'cuda'), + (2, 64, 128, 3, "swish", True, False, torch.float32, 'cuda'), + (2, 128, 128, 4, "swish", False, False, torch.float32, 'cuda'), + (2, 2, 128, 4, "swish", True, True, torch.float32, 'cuda'), # T_prefill < W + (2, 2, 128, 4, "swish", True, True, torch.float32, 'triton'), + (2, 3, 128, 4, "swish", True, True, torch.float32, 'triton'), + (2, 4, 128, 4, "swish", True, True, torch.float32, 'triton'), + (2, 2, 128, 3, "swish", True, True, torch.float32, 'triton'), + ] + ], +) +@torch.no_grad +def test_conv_with_cache_prefill_fwd( + B: int, + T: int, + D: int, + W: int, + activation: str, + has_bias: bool, + has_residual: bool, + dtype: torch.dtype, + backend: str, +): + if causal_conv1d_fn is None and backend == 'cuda': + pytest.skip("causal_conv1d is not installed for CUDA backend") + torch.manual_seed(42) + + x = torch.randn(B, T, D).to(device, dtype) + residual = torch.randn(B, T, D).to(device, dtype) if has_residual else None + + conv = ShortConvolution( + hidden_size=D, + kernel_size=W, + bias=has_bias, + activation=activation, + backend=backend, + device=device, + dtype=dtype, + ) + + cache = torch.randn(B, D, W - 1).to(device, dtype) + + ref = causal_conv1d_ref_torch( + x=x.transpose(1, 2), # (B, D, T) + weight=rearrange(conv.weight, "d 1 w -> d w"), + bias=conv.bias, + initial_state=cache, # (B, D, W-1) + activation=activation, + ).transpose(1, 2) # (B, T, D) + if has_residual: + ref += residual + + zero_padding = torch.zeros(B, D, 1).to(device, dtype) + tri_cache = torch.cat([zero_padding, cache], dim=-1) # (B, D, W) + tri, cache_out = conv(x, residual=residual, cache=tri_cache.clone(), output_final_state=True) + + assert_close("y", ref, tri, 1e-3) + for p in range(1, W): + if p <= T: + expected = x[:, -p, :] + else: + expected = tri_cache[:, :, -(p - T)] + torch.testing.assert_close( + cache_out[:, :, -p], + expected, + atol=1e-3, rtol=1e-3, + ) + + +@pytest.mark.parametrize( + ('N', 'T', 'D', 'W', 'activation', 'has_bias', 'has_residual', 'dtype', 'backend'), + [ + pytest.param( + *test, + id="N{0}_T{1}_D{2}_W{3}_activation{4}_has_bias{5}_has_residual{6}_{7}_{8}".format(*test), + ) + for test in [ + (3, 128, 64, 4, "swish", True, True, torch.float32, 'triton'), + (4, 256, 128, 3, None, False, True, torch.float32, 'triton'), + (2, 64, 128, 4, "swish", True, False, torch.float16, 'cuda'), + (3, 200, 64, 3, None, False, False, torch.float16, 'cuda'), + (2, 3, 64, 4, "swish", True, True, torch.float32, 'triton'), # T < W + (2, 3, 64, 3, None, False, True, torch.float32, 'cuda'), # T < W + ] + ], +) +@torch.no_grad +def test_conv_varlen_with_cache_prefill_fwd( + N: int, + T: int, + D: int, + W: int, + activation: str, + has_bias: bool, + has_residual: bool, + dtype: torch.dtype, + backend: str, +): + if causal_conv1d_fn is None and backend == 'cuda': + pytest.skip("causal_conv1d is not installed for CUDA backend") + torch.manual_seed(42) + + min_len_each = max(1, T // N) + lengths = [min_len_each] * N + lengths[-1] += T % N + assert all(length >= 1 for length in lengths), "all lengths must >= 1" + cu_seqlens = torch.tensor([0] + torch.cumsum(torch.tensor(lengths), 0).tolist(), + device=device, dtype=torch.int32) + + x = torch.randn(1, T, D).to(device, dtype) + residual = torch.randn(1, T, D).to(device, dtype) if has_residual else None + + conv = ShortConvolution( + hidden_size=D, + kernel_size=W, + bias=has_bias, + activation=activation, + backend=backend, + device=device, + dtype=dtype, + ) + + cache = torch.randn(N, D, W - 1).to(device, dtype) + ref_list = [] + for i, (bos, eos) in enumerate(zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False)): + xi = x[:, bos:eos, :].transpose(1, 2) # (1, D, l) + ci = cache[i:i + 1] # (1, D, W-1) + refi = causal_conv1d_ref_torch( + x=xi, + weight=rearrange(conv.weight, "d 1 w -> d w"), + bias=conv.bias, + initial_state=ci, + activation=activation, + ).transpose(1, 2) # (1, l, D) + if has_residual: + refi += residual[:, bos:eos, :] + ref_list.append(refi) + ref = torch.cat(ref_list, dim=1) # (1, T, D) + + zero_pad = torch.zeros(N, D, 1, device=device, dtype=dtype) + tri_cache = torch.cat([zero_pad, cache], dim=-1) # (N, D, W) + tri, cache_out = conv(x, + residual=residual, + cache=tri_cache.clone(), + cu_seqlens=cu_seqlens, + output_final_state=True) + + assert_close("varlen y", ref, tri, 1e-3) + + for i, (bos, eos) in enumerate(zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False)): + length = eos - bos + for p in range(1, W): + if p <= length: + expected = x[0, eos - p, :] + else: + expected = tri_cache[i, :, -(p - length)] + torch.testing.assert_close( + cache_out[i, :, -p], + expected, + atol=1e-3, + rtol=1e-3, + ) + + +@pytest.mark.parametrize( + ('B', 'D', 'W', 'has_bias', 'has_residual', 'activation', 'dtype', 'backend'), + [ + pytest.param(*test, id="B{0}_D{1}_W{2}_has_bias{3}_has_residual{4}_activation{5}_{6}_{7}".format(*test)) + for test in [ + (2, 128, 3, True, True, "swish", torch.float32, 'triton'), + (2, 128, 4, False, True, "swish", torch.float32, 'triton'), + (2, 128, 3, True, False, "swish", torch.float32, 'triton'), + (2, 128, 4, False, False, "swish", torch.float32, 'triton'), + (2, 128, 3, True, True, "swish", torch.float32, 'cuda'), + (2, 128, 4, False, True, "swish", torch.float32, 'cuda'), + (2, 128, 3, True, False, "swish", torch.float32, 'cuda'), + (2, 128, 4, False, False, "swish", torch.float32, 'cuda'), + (2, 128, 4, False, False, None, torch.float32, 'cuda'), + (2, 128, 4, False, False, None, torch.float32, 'triton'), + ] + ], +) +@torch.no_grad +def test_conv_decoding_with_cache( + B: int, + D: int, + W: int, + activation: str, + has_bias: bool, + has_residual: bool, + dtype: torch.dtype, + backend: str, +): + if causal_conv1d_fn is None and backend == 'cuda': + pytest.skip("causal_conv1d is not installed for CUDA backend") + torch.manual_seed(42) + + x = torch.randn(B, 1, D).to(device, dtype) # (B, 1, D) + residual = x.clone() if has_residual else None + + conv = ShortConvolution( + hidden_size=D, + kernel_size=W, + bias=has_bias, + activation=activation, + backend=backend, + device=device, + dtype=dtype, + ) + + state = torch.randn(B, D, W).to(device, dtype) + + # reference + ref = causal_conv1d_update_ref_torch( + x.squeeze(1), # (B, D) + conv_state=state.clone(), + weight=rearrange(conv.weight, "d 1 w -> d w"), + bias=conv.bias, + activation=activation, + ).unsqueeze(1) # (B, 1, D) + if has_residual: + ref += residual + + # ShortConvolution step + with torch.no_grad(): + y, _ = conv.step(x, residual, state.clone()) + + assert_close("y", ref, y, 1e-3) + + +@pytest.mark.parametrize( + ('B', 'T', 'D', 'W', 'has_bias', 'has_residual', 'activation', 'dtype'), + [ + pytest.param(*test, id="B{0}_T{1}_D{2}_W{3}_has_bias{4}_has_residual{5}_activation{6}_{7}".format(*test)) + for test in [ + (2, 64, 128, 3, True, True, "swish", torch.float32), + (2, 128, 128, 4, False, True, "swish", torch.float32), + (2, 64, 128, 3, True, False, "swish", torch.float32), + (2, 128, 128, 4, False, False, "swish", torch.float32), + ] + ], +) +@torch.no_grad +def test_mixed_backend( + B: int, + T: int, + D: int, + W: int, + has_bias: bool, + has_residual: bool, + activation: str, + dtype: torch.dtype, +): + torch.manual_seed(42) + T_decode = 1 + x = torch.randn(B, T + T_decode, D, device=device, dtype=dtype) + residual = torch.randn_like(x) if has_residual else None + + conv = ShortConvolution( + hidden_size=D, + kernel_size=W, + bias=has_bias, + activation=activation, + backend="cuda", + device=device, + dtype=dtype, + ) + + cache = torch.randn(B, D, W-1, device=device, dtype=dtype) + y_cuda_prefill, final_state = conv( + x[:, :T], + residual=residual[:, :T] if has_residual else None, + cache=cache, + output_final_state=True, + ) + + conv.backend = "triton" + y_triton_decode, _ = conv( + x[:, T:], + residual=residual[:, T:] if has_residual else None, + cache=final_state, + output_final_state=True, + ) + + conv.backend = "triton" + cache = torch.cat((torch.zeros_like(cache[..., :1]), cache), -1) + y_triton_full, _ = conv(x, residual=residual, cache=cache) + + y_mixed = torch.cat([y_cuda_prefill, y_triton_decode], dim=1) + assert_close("cuda→triton vs triton", y_mixed, y_triton_full, 1e-3) + + conv.backend = "triton" + y_triton_prefill, final_state = conv( + x[:, :T], + residual=residual[:, :T] if has_residual else None, + cache=cache, + output_final_state=True, + ) + + conv.backend = "cuda" + y_cuda_decode, _ = conv( + x[:, T:], + residual=residual[:, T:] if has_residual else None, + cache=final_state, + output_final_state=True, + ) + + y_mixed2 = torch.cat([y_triton_prefill, y_cuda_decode], dim=1) + assert_close("triton→cuda vs triton", y_mixed2, y_triton_full, 1e-3) + + +@pytest.mark.parametrize( + ('B', 'T', 'D', 'W', 'has_bias', 'has_residual', 'activation', 'dtype'), + [ + pytest.param(*test, id="B{0}_T{1}_D{2}_W{3}_has_bias{4}_has_residual{5}_activation{6}_{7}".format(*test)) + for test in [ + (2, 64, 100, 3, True, True, "swish", torch.float32), + (2, 128, 128, 4, True, True, "swish", torch.float32), + (3, 128, 128, 4, True, True, "swish", torch.float32), + (3, 128, 256, 4, True, True, "swish", torch.float32), + (3, 128, 512, 4, True, True, "swish", torch.float32), + (2, 128, 1024, 4, True, True, "swish", torch.float32), + (2, 128, 2048, 3, True, True, "swish", torch.float32), + (2, 128, 4096, 4, True, True, "swish", torch.float32), + (2, 128, 8192, 4, True, True, "swish", torch.float32), + ] + ], +) +def test_conv_cache_backward( + B: int, + T: int, + D: int, + W: int, + has_bias: bool, + has_residual: bool, + activation: str, + dtype: torch.dtype, +): + torch.manual_seed(42) + + x = torch.randn(B, T, D, device=device, dtype=dtype, requires_grad=True) + weight = torch.randn(D, W, device=device, dtype=dtype, requires_grad=True) + bias = torch.randn(D, device=device, dtype=dtype, requires_grad=True) if has_bias else None + residual = torch.randn(B, T, D, device=device, dtype=dtype, requires_grad=True) if has_residual else None + cache = torch.randn(B, D, W - 1, device=device, dtype=dtype, requires_grad=True) + + def ref_func(x, weight, bias, residual, cache): + out, cache_out = causal_conv1d_ref_torch( + x.transpose(1, 2), + weight, + bias, + initial_state=cache, + output_final_state=True, + activation=activation, + ) + out = out.transpose(1, 2) + if residual is not None: + out += residual + return out, cache_out + + def triton_func(x, weight, bias, residual, cache): + zero_padding = torch.zeros(B, D, 1, device=device, dtype=dtype) + triton_cache = torch.cat([zero_padding, cache], dim=-1).contiguous() + tri, cache_out_triton = causal_conv1d( + x, + weight=weight, + bias=bias, + residual=residual, + initial_state=triton_cache, + output_final_state=True, + activation=activation, + ) + cache_out_triton = cache_out_triton[..., 1:].clone() # [B, D, W-1] + return tri, cache_out_triton + + d_tri = torch.randn_like(x) + d_cache_out = torch.randn_like(cache) + + def get_grads(func, *inputs): + out, cache_out = func(*inputs) + loss = (out * d_tri).sum() + (cache_out * d_cache_out).sum() + grads = torch.autograd.grad( + loss, + inputs, + retain_graph=True, + create_graph=False, + ) + return grads + + inputs = (x, weight, bias, residual, cache) + grads_ref = get_grads(ref_func, *inputs) + grads_tri = get_grads(triton_func, *inputs) + + names = ["x", "weight", "bias", "residual", "cache"] + for name, g_ref, g_tri in zip(names, grads_ref, grads_tri, strict=False): + assert_close(name, g_ref, g_tri, ratio=1e-3) diff --git a/code/flash-linear-attention/tests/modules/test_cross_entropy.py b/code/flash-linear-attention/tests/modules/test_cross_entropy.py new file mode 100644 index 0000000000000000000000000000000000000000..71966bd7ce5e78f960b5223f6830d0d4a1364938 --- /dev/null +++ b/code/flash-linear-attention/tests/modules/test_cross_entropy.py @@ -0,0 +1,81 @@ + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss +from fla.utils import assert_close, device, device_platform + + +@pytest.mark.parametrize("B", [2]) +@pytest.mark.parametrize("T", [512, 1024]) +@pytest.mark.parametrize("D", [1024, 2048]) +@pytest.mark.parametrize("V", [32000, 100000]) +@pytest.mark.parametrize("reduction", ['mean']) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.skipif( + device_platform == 'intel', + reason="Intel Triton Failure", +) +def test_fused_cross_entropy(B: int, T: int, D: int, V: int, reduction: str, dtype: torch.dtype): + torch.manual_seed(42) + logits = torch.randn(B * T, V).to(device).to(dtype=dtype).requires_grad_() + target = torch.randint(0, V, (B, T)).to(device) + target = torch.cat((target[..., 1:], torch.full_like(target[..., :1], -100)), -1) + target = target.flatten() + + ref = nn.CrossEntropyLoss(reduction=reduction)(logits, target).to(dtype=dtype) + do = torch.randn_like(ref).to(device).to(dtype=dtype) + + ref.backward(do) + ref_d, logits.grad = logits.grad.clone(), None + + tri = FusedCrossEntropyLoss(reduction=reduction)(logits, target).to(dtype=dtype) + tri.backward(do) + tri_d, logits.grad = logits.grad.clone(), None + + assert_close(" o", ref, tri, ratio=1e-2) + assert_close("dl", ref_d, tri_d, ratio=1e-2) + + +@pytest.mark.parametrize("B", [2]) +@pytest.mark.parametrize("T", [512, 1024]) +@pytest.mark.parametrize("D", [1024, 2048]) +@pytest.mark.parametrize("V", [32000, 100000]) +@pytest.mark.parametrize("scale", [1., 0.5]) +@pytest.mark.parametrize("reduction", ['mean']) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.skipif( + device_platform == 'intel', + reason="Intel Triton Failure", +) +def test_fused_linear_cross_entropy(B: int, T: int, D: int, V: int, scale: float, reduction: str, dtype: torch.dtype): + torch.manual_seed(42) + + x = torch.randn(B * T, D).to(device).to(dtype=dtype).requires_grad_() + target = torch.randint(0, V, (B, T)).to(device) + target = torch.cat((target[..., 1:], torch.full_like(target[..., :1], -100)), -1) + target = target.flatten() + weight = torch.randn(V, D).to(device).to(dtype=dtype).requires_grad_() + bias = torch.randn(V).to(device).to(dtype=dtype).requires_grad_() + + logits = F.linear(x, weight, bias) + ref = FusedCrossEntropyLoss(logit_scale=scale, reduction=reduction)(logits, target) + do = torch.randn_like(ref).to(device).to(dtype=dtype) + + ref.backward(do) + ref_dx, x.grad = x.grad.clone(), None + ref_dw, weight.grad = weight.grad.clone(), None + ref_db, bias.grad = bias.grad.clone(), None + + tri = FusedLinearCrossEntropyLoss(logit_scale=scale, reduction=reduction)(x, target, weight, bias) + tri.backward(do) + tri_dx, x.grad = x.grad.clone(), None + tri_dw, weight.grad = weight.grad.clone(), None + tri_db, bias.grad = bias.grad.clone(), None + + assert_close(" o", ref, tri, ratio=1e-2) + assert_close("dx", ref_dx, tri_dx, ratio=1e-2) + assert_close("dw", ref_dw, tri_dw, ratio=1e-2) + assert_close("db", ref_db, tri_db, ratio=1e-2) diff --git a/code/flash-linear-attention/tests/modules/test_grpo.py b/code/flash-linear-attention/tests/modules/test_grpo.py new file mode 100644 index 0000000000000000000000000000000000000000..965065860aa09cc5a32e06008540f8d5c5e90b5e --- /dev/null +++ b/code/flash-linear-attention/tests/modules/test_grpo.py @@ -0,0 +1,56 @@ + +import pytest +import torch + +from fla.modules.grpo import fused_grpo_loss, grpo_loss_torch +from fla.utils import assert_close, device, device_torch_lib, is_nvidia_hopper + + +@pytest.mark.parametrize("B", [2]) +@pytest.mark.parametrize("T", [16, 1024, 4096]) +@pytest.mark.parametrize("V", [32000, 65536, 131072]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("inplace", [True, False]) +@pytest.mark.parametrize("repeat", [100]) +def test_fused_grpos(B: int, T: int, V: int, dtype: torch.dtype, inplace: bool, repeat: int): + device_torch_lib.manual_seed(42) + for i in range(repeat): + if not is_nvidia_hopper and T == 4096: + pytest.skip("Skip test for T=4096 on Intel Alchemist") + + def get_random_ref_log_probs(logits, input_ids): + with torch.inference_mode(): + logits = logits[:, :-1] + per_token_logps = [] + for logits_row, input_ids_row in zip(logits, input_ids[:, -logits.size(1):], strict=False): + log_probs = torch.randn_like(logits_row).log_softmax(dim=-1) + token_log_prob = torch.gather(log_probs, dim=1, index=input_ids_row.unsqueeze(1)).squeeze(1) + per_token_logps.append(token_log_prob) + device_torch_lib.empty_cache() + return torch.stack(per_token_logps) + + logits = torch.randn(B, T + 1, V, device=device, dtype=dtype) + logits.requires_grad_(True) + advantages = torch.randn(B, device=device, dtype=torch.float32) + input_ids = torch.randint(0, V-1, (B, T + 64), device=device) + ref_logp = get_random_ref_log_probs(logits, input_ids) + beta = 0.04 + completion_mask = torch.ones(B, T, dtype=torch.int32, device=device) + completion_mask[::2, T//3: T//2] = 0 + save_kl = True + + gold_logits = logits.detach().clone().float() + gold_logits.requires_grad_(True) + gold_ref_logp = ref_logp.clone().float() + device_torch_lib.empty_cache() + y1 = fused_grpo_loss(logits, ref_logp, input_ids, advantages, beta, completion_mask, save_kl=save_kl, inplace=inplace) + y2 = grpo_loss_torch(gold_logits, gold_ref_logp, input_ids, advantages, beta, completion_mask, save_kl) + if save_kl: + y1, kl2 = y1 + y2, kl3 = y2 + assert (kl2-kl3).abs().max() < 1e-3 + dy = torch.randn_like(y1) * 10 + y1.backward(dy) + y2.backward(dy.float()) + assert (y1-y2).abs().max() < 1e-3 + assert_close(" dlogits", gold_logits.grad, logits.grad, 3e-3) diff --git a/code/flash-linear-attention/tests/modules/test_kl_div.py b/code/flash-linear-attention/tests/modules/test_kl_div.py new file mode 100644 index 0000000000000000000000000000000000000000..3b119b5c3742456798a13465bd81ed3fc54c2b24 --- /dev/null +++ b/code/flash-linear-attention/tests/modules/test_kl_div.py @@ -0,0 +1,44 @@ + +import pytest +import torch +import torch.nn.functional as F + +from fla.modules import FusedKLDivLoss +from fla.utils import assert_close, device, device_platform + + +@pytest.mark.parametrize("B", [2]) +@pytest.mark.parametrize("T", [16, 32]) +@pytest.mark.parametrize("D", [1024, 2048]) +@pytest.mark.parametrize("V", [32000, 100000]) +@pytest.mark.parametrize("reduction", ["batchmean"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +@pytest.mark.skipif( + device_platform == 'intel', + reason="Intel Triton Failure", +) +def test_fused(B: int, T: int, D: int, V: int, reduction: str, dtype: torch.dtype): + torch.manual_seed(42) + x = torch.randn(B * T, D).to(device).to(dtype=dtype).requires_grad_() + x_weight = torch.randn(V, D).to(device).to(dtype=dtype).requires_grad_() + target_x = torch.randn(B * T, D).to(device).to(dtype=dtype) + target_weight = torch.randn(V, D).to(device).to(dtype=dtype) + + ref = F.kl_div( + F.linear(x, x_weight).log_softmax(-1), + F.linear(target_x, target_weight).softmax(-1), + reduction=reduction, + ).to(dtype) + do = torch.randn_like(ref).to(device) + ref.backward(do) + ref_dx, x.grad = x.grad.clone(), None + ref_dw, x_weight.grad = x_weight.grad.clone(), None + + tri = FusedKLDivLoss(reduction)(x, target_x, x_weight, target_weight).to(dtype=dtype) + tri.backward(do) + tri_dx, x.grad = x.grad.clone(), None + tri_dw, x_weight.grad = x_weight.grad.clone(), None + + assert_close(" o", ref, tri, 1e-2) + assert_close(" dx", ref_dx, tri_dx, 1e-2) + assert_close(" dw", ref_dw, tri_dw, 1e-2) diff --git a/code/flash-linear-attention/tests/modules/test_l2norm.py b/code/flash-linear-attention/tests/modules/test_l2norm.py new file mode 100644 index 0000000000000000000000000000000000000000..18f882e57796c6d4a45a6d7bd722c10bfe7532c2 --- /dev/null +++ b/code/flash-linear-attention/tests/modules/test_l2norm.py @@ -0,0 +1,37 @@ + +import pytest +import torch +import torch.nn.functional as F + +from fla.modules.l2norm import l2_norm +from fla.utils import assert_close, device + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (1, 63, 1, 60, torch.float), + (2, 500, 4, 64, torch.float), + (2, 1000, 2, 100, torch.float), + (3, 1024, 4, 128, torch.float), + (4, 1024, 5, 1024, torch.float16), + (4, 1024, 5, 1024, torch.bfloat16), + (5, 1024, 6, 2048, torch.float16), + (5, 1024, 6, 2048, torch.bfloat16), + ] + ], +) +def test_l2norm(B: int, T: int, H: int, D: int, dtype: torch.dtype): + torch.manual_seed(42) + x = torch.randn(B, T, H, D, dtype=dtype).to(device).requires_grad_(True) + x = x * 0.5 + 0.3 + + ref = F.normalize(x, dim=-1, p=2) + tri = l2_norm(x) + ref_dx = torch.autograd.grad(ref.sum(), x)[0] + tri_dx = torch.autograd.grad(tri.sum(), x)[0] + + assert_close('y', ref, tri, 0.005) + assert_close('dx', ref_dx, tri_dx, 0.005) diff --git a/code/flash-linear-attention/tests/modules/test_l2warp.py b/code/flash-linear-attention/tests/modules/test_l2warp.py new file mode 100644 index 0000000000000000000000000000000000000000..b0ce683e9d4bc30061a999e8014f5139b79db85b --- /dev/null +++ b/code/flash-linear-attention/tests/modules/test_l2warp.py @@ -0,0 +1,70 @@ + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from fla.modules import FusedLinearCrossEntropyLoss +from fla.modules.l2warp import l2_warp as standalone_l2_warp +from fla.utils import assert_close, device, is_intel_alchemist + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("B", [4, 8]) +@pytest.mark.parametrize("T", [1024]) +@pytest.mark.parametrize("H", [256]) +@pytest.mark.parametrize("V", [2000]) +@pytest.mark.parametrize("l2_penalty_factor", [1e-4, 1]) +@pytest.mark.skipif( + is_intel_alchemist is True, + reason="Intel Triton Failure", +) +def test_fused_linear_cross_entropy_l2_warp( + B: int, + T: int, + H: int, + V: int, + l2_penalty_factor: float, + dtype: torch.dtype, +): + torch.manual_seed(42) + + lm_head = nn.Linear(H, V, bias=True, device=device, dtype=dtype) + x = torch.randn(B, T, H, device=device, dtype=dtype, requires_grad=True) + labels = torch.randint(0, V, (B, T), device=device) + + ignore_index = -100 + shift_labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], ignore_index)), 1) + + ref_criterion = nn.CrossEntropyLoss() + + ref_logits = F.linear(x.view(-1, H), lm_head.weight, lm_head.bias) + ref_loss_ce = ref_criterion(ref_logits.view(B * T, V), shift_labels.view(-1)) + ref_loss = standalone_l2_warp(ref_loss_ce, ref_logits.view(B, T, V), l2_penalty_factor) + + ref_loss.backward() + ref_x_grad = x.grad.clone() + ref_w_grad = lm_head.weight.grad.clone() + ref_b_grad = lm_head.bias.grad.clone() + + x.grad = None + lm_head.zero_grad() + + fused_criterion = FusedLinearCrossEntropyLoss( + l2_penalty_factor=l2_penalty_factor, + use_l2warp=True, # Make sure to enable it + ) + + fused_loss = fused_criterion(x, shift_labels, lm_head.weight, lm_head.bias) + + fused_loss.backward() + fused_x_grad = x.grad.clone() + fused_w_grad = lm_head.weight.grad.clone() + fused_b_grad = lm_head.bias.grad.clone() + + ratio = 4e-3 if dtype == torch.bfloat16 else 1e-3 + + assert_close("Loss", ref_loss, fused_loss, ratio) + assert_close("dx", ref_x_grad, fused_x_grad, ratio) + assert_close("dw", ref_w_grad, fused_w_grad, ratio) + assert_close("db", ref_b_grad, fused_b_grad, ratio) diff --git a/code/flash-linear-attention/tests/modules/test_layernorm.py b/code/flash-linear-attention/tests/modules/test_layernorm.py new file mode 100644 index 0000000000000000000000000000000000000000..971c9abf4a62dc8ba970ece291cd8608fb14e640 --- /dev/null +++ b/code/flash-linear-attention/tests/modules/test_layernorm.py @@ -0,0 +1,217 @@ + +import pytest +import torch +import torch.nn as nn +from einops import rearrange +from transformers.models.llama.modeling_llama import LlamaRMSNorm + +from fla.modules import GroupNorm, GroupNormLinear, LayerNorm, LayerNormLinear, RMSNorm, RMSNormLinear +from fla.modules.layernorm import GroupNormRef +from fla.utils import assert_close, device + + +@pytest.mark.parametrize("B", [2]) +@pytest.mark.parametrize("H", [2]) +@pytest.mark.parametrize("T", [512]) +@pytest.mark.parametrize("D", [50, 64, 128]) +@pytest.mark.parametrize("elementwise_affine", [False, True]) +@pytest.mark.parametrize("bias", [False, True]) +def test_layernorm(B: int, H: int, T: int, D: int, elementwise_affine: bool, bias: bool): + x = torch.randn(B, H, T, D).to(device).requires_grad_(True) + ref = nn.LayerNorm(D, elementwise_affine=elementwise_affine, bias=bias).to(device) + tri = LayerNorm(D, elementwise_affine=elementwise_affine, bias=bias).to(device) + if ref.weight is not None: + nn.init.normal_(ref.weight) + tri.weight.data.copy_(ref.weight.data) + if ref.bias is not None: + nn.init.normal_(ref.bias) + tri.bias.data.copy_(ref.bias.data) + + ref_y = ref(x) + tri_y = tri(x) + ref_dx = torch.autograd.grad(ref(x).sum(), x)[0] + tri_dx = torch.autograd.grad(tri(x).sum(), x)[0] + + if ref.weight is not None: + ref_dw = torch.autograd.grad(ref(x).sum(), ref.weight)[0] + tri_dw = torch.autograd.grad(tri(x).sum(), tri.weight)[0] + if ref.bias is not None: + ref_db = torch.autograd.grad(ref(x).sum(), ref.bias)[0] + tri_db = torch.autograd.grad(tri(x).sum(), tri.bias)[0] + + assert_close(' y', ref_y, tri_y, 1e-3) + assert_close('dx', ref_dx, tri_dx, 1e-3) + if ref.weight is not None: + assert_close('dw', ref_dw, tri_dw, 1e-3) + if ref.bias is not None: + assert_close('db', ref_db, tri_db, 1e-3) + + +@pytest.mark.parametrize("B", [2]) +@pytest.mark.parametrize("T", [512]) +@pytest.mark.parametrize("D", [64, 128, 512, 1024, 2048]) +@pytest.mark.parametrize("G", [1, 4]) +@pytest.mark.parametrize("is_rms_norm", [True, False]) +def test_groupnorm(B: int, T: int, D: int, G: int, is_rms_norm: bool): + torch.manual_seed(42) + x = torch.randn(B, T, D).to(device).requires_grad_(True) + if is_rms_norm: + ref = GroupNormRef(num_groups=G, hidden_size=D, bias=True, is_rms_norm=True).to(device) + else: + ref = nn.GroupNorm(G, D).to(device) + tri = GroupNorm(G, D, bias=True, is_rms_norm=is_rms_norm).to(device) + nn.init.normal_(ref.weight) + nn.init.normal_(ref.bias) + tri.weight.data.copy_(ref.weight.data) + tri.bias.data.copy_(ref.bias.data) + ref = ref.to(dtype=torch.float32) + + ref_x = rearrange(x, 'b t d -> (b t) d').to(dtype=torch.float32) + ref_y = rearrange(ref(ref_x), '(b t) d -> b t d', b=B) + tri_y = tri(x) + ref_dx = torch.autograd.grad(ref(ref_x).sum(), x)[0] + tri_dx = torch.autograd.grad(tri(x).sum(), x)[0] + ref_dw = torch.autograd.grad(ref(ref_x).sum(), ref.weight)[0] + tri_dw = torch.autograd.grad(tri(x).sum(), tri.weight)[0] + ref_db = torch.autograd.grad(ref(ref_x).sum(), ref.bias)[0] + tri_db = torch.autograd.grad(tri(x).sum(), tri.bias)[0] + + assert_close(' y', ref_y, tri_y, 1e-3) + assert_close('dx', ref_dx, tri_dx, 1e-3) + assert_close('dw', ref_dw, tri_dw, 1e-3) + assert_close('db', ref_db, tri_db, 1e-3) + + +@pytest.mark.parametrize("B", [2]) +@pytest.mark.parametrize("H", [2]) +@pytest.mark.parametrize("T", [512]) +@pytest.mark.parametrize("D", [50, 64, 128]) +def test_rmsnorm(B: int, H: int, T: int, D: int): + x = torch.randn(B, H, T, D).to(device).requires_grad_(True) + ref = LlamaRMSNorm(D, eps=0).to(device) + tri = RMSNorm(D, eps=0).to(device) + nn.init.normal_(ref.weight) + tri.weight.data.copy_(ref.weight.data) + + ref_y = ref(x) + tri_y = tri(x) + ref_dx = torch.autograd.grad(ref(x).sum(), x)[0] + tri_dx = torch.autograd.grad(tri(x).sum(), x)[0] + + ref_dw = torch.autograd.grad(ref(x).sum(), ref.weight)[0] + tri_dw = torch.autograd.grad(tri(x).sum(), tri.weight)[0] + + assert_close(' y', ref_y, tri_y, 1e-3) + assert_close('dx', ref_dx, tri_dx, 1e-3) + assert_close('dw', ref_dw, tri_dw, 1e-3) + + +@pytest.mark.parametrize("N", [1, 16, 128]) +@pytest.mark.parametrize("D", [50, 64, 128]) +def test_layernorm_linear(N: int, D: int): + torch.manual_seed(1) + x = torch.randn(N, D).to(device).requires_grad_(True) + ref = nn.Sequential(nn.LayerNorm(D, elementwise_affine=True, bias=True), nn.Linear(D, D)).to(device) + tri = LayerNormLinear(D, elementwise_affine=True, bias=True).to(device) + nn.init.normal_(ref[0].weight) + nn.init.normal_(ref[0].bias) + nn.init.normal_(ref[1].weight, mean=0.0, std=0.01) + nn.init.normal_(ref[1].bias, mean=0.0, std=0.01) + tri.weight.data.copy_(ref[0].weight.data) + tri.bias.data.copy_(ref[0].bias.data) + weight, bias = ref[1].weight.clone(), ref[1].bias.clone() + + ref_y = ref(x) + tri_y = tri(x, weight, bias) + ref_dx = torch.autograd.grad(ref(x).sum(), x)[0] + tri_dx = torch.autograd.grad(tri(x, weight, bias).sum(), x)[0] + ref_dw = torch.autograd.grad(ref(x).sum(), ref[0].weight)[0] + tri_dw = torch.autograd.grad(tri(x, weight, bias).sum(), tri.weight)[0] + ref_db = torch.autograd.grad(ref(x).sum(), ref[0].bias)[0] + tri_db = torch.autograd.grad(tri(x, weight, bias).sum(), tri.bias)[0] + ref_dlw = torch.autograd.grad(ref(x).sum(), ref[1].weight)[0] + tri_dlw = torch.autograd.grad(tri(x, weight, bias).sum(), weight)[0] + ref_dlb = torch.autograd.grad(ref(x).sum(), ref[1].bias)[0] + tri_dlb = torch.autograd.grad(tri(x, weight, bias).sum(), bias)[0] + + assert_close(' y', ref_y, tri_y, 1e-3) + assert_close(' dx', ref_dx, tri_dx, 1e-3) + assert_close(' dw', ref_dw, tri_dw, 1e-3) + assert_close(' db', ref_db, tri_db, 1e-3) + assert_close('dlw', ref_dlw, tri_dlw, 1e-3) + assert_close('dlb', ref_dlb, tri_dlb, 1e-3) + + +@pytest.mark.parametrize("N", [1, 16, 128]) +@pytest.mark.parametrize("D", [64, 128, 512]) +@pytest.mark.parametrize("G", [1, 4]) +@pytest.mark.parametrize("is_rms_norm", [True, False]) +def test_groupnorm_linear(N: int, D: int, G: int, is_rms_norm: bool): + torch.manual_seed(1) + x = torch.randn(N, D).to(device).requires_grad_(True) + if is_rms_norm: + ref = nn.Sequential( + GroupNormRef(num_groups=G, hidden_size=D, bias=True, is_rms_norm=True), + nn.Linear(D, D), + ).to(device) + else: + ref = nn.Sequential(nn.GroupNorm(G, D), nn.Linear(D, D)).to(device) + tri = GroupNormLinear(G, D, bias=True, is_rms_norm=is_rms_norm).to(device) + nn.init.normal_(ref[0].weight) + nn.init.normal_(ref[0].bias) + nn.init.normal_(ref[1].weight, mean=0.0, std=0.01) + nn.init.normal_(ref[1].bias, mean=0.0, std=0.01) + tri.weight.data.copy_(ref[0].weight.data) + tri.bias.data.copy_(ref[0].bias.data) + weight, bias = ref[1].weight.clone(), ref[1].bias.clone() + + ref_y = ref(x) + tri_y = tri(x, weight, bias) + ref_dx = torch.autograd.grad(ref(x).sum(), x)[0] + tri_dx = torch.autograd.grad(tri(x, weight, bias).sum(), x)[0] + ref_dw = torch.autograd.grad(ref(x).sum(), ref[0].weight)[0] + tri_dw = torch.autograd.grad(tri(x, weight, bias).sum(), tri.weight)[0] + ref_db = torch.autograd.grad(ref(x).sum(), ref[0].bias)[0] + tri_db = torch.autograd.grad(tri(x, weight, bias).sum(), tri.bias)[0] + ref_dlw = torch.autograd.grad(ref(x).sum(), ref[1].weight)[0] + tri_dlw = torch.autograd.grad(tri(x, weight, bias).sum(), weight)[0] + ref_dlb = torch.autograd.grad(ref(x).sum(), ref[1].bias)[0] + tri_dlb = torch.autograd.grad(tri(x, weight, bias).sum(), bias)[0] + + assert_close(' y', ref_y, tri_y, 1e-3) + assert_close(' dx', ref_dx, tri_dx, 1e-3) + assert_close(' dw', ref_dw, tri_dw, 1e-3) + assert_close(' db', ref_db, tri_db, 1e-3) + assert_close('dlw', ref_dlw, tri_dlw, 1e-3) + assert_close('dlb', ref_dlb, tri_dlb, 1e-3) + + +@pytest.mark.parametrize("N", [1, 16, 128]) +@pytest.mark.parametrize("D", [50, 64, 128]) +def test_rmsnorm_linear(N: int, D: int): + torch.manual_seed(1) + x = torch.randn(N, D).to(device).requires_grad_(True) + ref = nn.Sequential(LlamaRMSNorm(D, eps=0), nn.Linear(D, D)).to(device) + tri = RMSNormLinear(D, eps=0).to(device) + nn.init.normal_(ref[0].weight) + nn.init.normal_(ref[1].weight, mean=0.0, std=0.01) + nn.init.normal_(ref[1].bias, mean=0.0, std=0.01) + tri.weight.data.copy_(ref[0].weight.data) + weight, bias = ref[1].weight.clone(), ref[1].bias.clone() + + ref_y = ref(x) + tri_y = tri(x, weight, bias) + ref_dx = torch.autograd.grad(ref(x).sum(), x)[0] + tri_dx = torch.autograd.grad(tri(x, weight, bias).sum(), x)[0] + ref_dw = torch.autograd.grad(ref(x).sum(), ref[0].weight)[0] + tri_dw = torch.autograd.grad(tri(x, weight, bias).sum(), tri.weight)[0] + ref_dlw = torch.autograd.grad(ref(x).sum(), ref[1].weight)[0] + tri_dlw = torch.autograd.grad(tri(x, weight, bias).sum(), weight)[0] + ref_dlb = torch.autograd.grad(ref(x).sum(), ref[1].bias)[0] + tri_dlb = torch.autograd.grad(tri(x, weight, bias).sum(), bias)[0] + + assert_close(' y', ref_y, tri_y, 1e-3) + assert_close(' dx', ref_dx, tri_dx, 1e-3) + assert_close(' dw', ref_dw, tri_dw, 1e-3) + assert_close('dlw', ref_dlw, tri_dlw, 1e-3) + assert_close('dlb', ref_dlb, tri_dlb, 1e-3) diff --git a/code/flash-linear-attention/tests/modules/test_layernorm_gated.py b/code/flash-linear-attention/tests/modules/test_layernorm_gated.py new file mode 100644 index 0000000000000000000000000000000000000000..13ef0b8e1646015a17e3d6dacb49b9e7e4dac68f --- /dev/null +++ b/code/flash-linear-attention/tests/modules/test_layernorm_gated.py @@ -0,0 +1,92 @@ + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from fla.modules import FusedLayerNormGated, FusedRMSNormGated +from fla.utils import assert_close, device + + +@pytest.mark.parametrize( + ('B', 'H', 'T', 'D', 'elementwise_affine', 'activation', 'bias'), + [ + pytest.param(*test, id=f"B{test[0]}_H{test[1]}_T{test[2]}_D{test[3]}_affine{test[4]}_{test[5]}_bias{test[6]}") + for test in [ + (2, 2, 1, 64, False, "silu", False), + (2, 2, 512, 128, True, "silu", True), + (2, 2, 2048, 1200, True, "sigmoid", False), + (2, 2, 50, 50, False, "sigmoid", False), + ] + ], +) +def test_layernorm_gated(B: int, H: int, T: int, D: int, elementwise_affine: bool, activation: str, bias: bool): + torch.manual_seed(42) + x = torch.randn(B, H, T, D).to(device).requires_grad_(True) + g = torch.randn(B, H, T, D).to(device).requires_grad_(True) + + ref = nn.LayerNorm(D, elementwise_affine=elementwise_affine, bias=bias).to(device) + tri = FusedLayerNormGated(D, elementwise_affine=elementwise_affine, bias=bias, activation=activation).to(device) + if ref.weight is not None: + nn.init.normal_(ref.weight) + tri.weight.data.copy_(ref.weight.data) + if ref.bias is not None: + nn.init.normal_(ref.bias) + tri.bias.data.copy_(ref.bias.data) + + act_fn = F.silu if activation == "silu" else F.sigmoid + ref_y = ref(x) * act_fn(g) + tri_y = tri(x, g) + ref_dx, ref_dg = torch.autograd.grad((ref(x) * act_fn(g)).sum(), (x, g)) + tri_dx, tri_dg = torch.autograd.grad(tri_y.sum(), (x, g)) + + if ref.weight is not None: + ref_dw = torch.autograd.grad((ref(x) * act_fn(g)).sum(), ref.weight)[0] + tri_dw = torch.autograd.grad(tri(x, g).sum(), tri.weight)[0] + if ref.bias is not None: + ref_db = torch.autograd.grad((ref(x) * act_fn(g)).sum(), ref.bias)[0] + tri_db = torch.autograd.grad(tri(x, g).sum(), tri.bias)[0] + + assert_close(' y', ref_y, tri_y, 1e-3) + assert_close('dx', ref_dx, tri_dx, 1e-3) + assert_close('dg', ref_dg, tri_dg, 1e-3) + if ref.weight is not None: + assert_close('dw', ref_dw, tri_dw, 1e-3) + if ref.bias is not None: + assert_close('db', ref_db, tri_db, 1e-3) + + +@pytest.mark.parametrize( + ('B', 'H', 'T', 'D', 'activation'), + [ + pytest.param(*test, id=f"B{test[0]}_H{test[1]}_T{test[2]}_D{test[3]}_{test[4]}") + for test in [ + (2, 2, 1, 64, "silu"), + (2, 2, 512, 128, "sigmoid"), + (2, 2, 2048, 1200, "silu"), + (2, 2, 50, 50, "sigmoid"), + ] + ], +) +def test_rmsnorm_gated(B: int, H: int, T: int, D: int, activation: str): + torch.manual_seed(42) + x = torch.randn(B, H, T, D).to(device).requires_grad_(True) + g = torch.randn(B, H, T, D).to(device).requires_grad_(True) + ref = nn.RMSNorm(D, eps=0).to(device) + tri = FusedRMSNormGated(D, eps=0, activation=activation).to(device) + nn.init.normal_(ref.weight) + tri.weight.data.copy_(ref.weight.data) + + act_fn = F.silu if activation == "silu" else F.sigmoid + ref_y = ref(x) * act_fn(g) + tri_y = tri(x, g) + ref_dx, ref_dg = torch.autograd.grad((ref(x) * act_fn(g)).sum(), (x, g)) + tri_dx, tri_dg = torch.autograd.grad(tri_y.sum(), (x, g)) + + ref_dw = torch.autograd.grad((ref(x) * act_fn(g)).sum(), ref.weight)[0] + tri_dw = torch.autograd.grad(tri(x, g).sum(), tri.weight)[0] + + assert_close(' y', ref_y, tri_y, 1e-3) + assert_close('dx', ref_dx, tri_dx, 1e-3) + assert_close('dg', ref_dg, tri_dg, 1e-3) + assert_close('dw', ref_dw, tri_dw, 1e-3) diff --git a/code/flash-linear-attention/tests/modules/test_rotary.py b/code/flash-linear-attention/tests/modules/test_rotary.py new file mode 100644 index 0000000000000000000000000000000000000000..fd95cf537caf6c28fb19d70a2bdce94ca744b8a5 --- /dev/null +++ b/code/flash-linear-attention/tests/modules/test_rotary.py @@ -0,0 +1,122 @@ + +import pytest +import torch + +from fla.modules.rotary import RotaryEmbedding, rotary_embedding_ref +from fla.utils import assert_close, device + + +@pytest.mark.parametrize("B", [2]) +@pytest.mark.parametrize("T", [2048, 4096]) +@pytest.mark.parametrize("H", [4]) +@pytest.mark.parametrize("G", [1, 4]) +@pytest.mark.parametrize("D", [128, 256]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +def test_rotary(B: int, T: int, H: int, G: int, D: int, dtype: torch.dtype): + torch.manual_seed(42) + q = torch.randn(B, T, H, D).to(device).to(dtype=dtype).requires_grad_() + k = torch.randn(B, T, H//G, D).to(device).to(dtype=dtype).requires_grad_() + rotary = RotaryEmbedding(D).to(device) + + tri_q, tri_k = rotary(q, k) + tri_dq = torch.autograd.grad(tri_q.sum(), q, retain_graph=True)[0] + tri_dk = torch.autograd.grad(tri_k.sum(), k, retain_graph=True)[0] + + ref_q = rotary_embedding_ref(q.float(), rotary._cos_cached, rotary._sin_cached).to(dtype=dtype) + ref_k = rotary_embedding_ref(k.float(), rotary._cos_cached, rotary._sin_cached).to(dtype=dtype) + ref_dq = torch.autograd.grad(ref_q.sum(), q, retain_graph=True)[0] + ref_dk = torch.autograd.grad(ref_k.sum(), k, retain_graph=True)[0] + + assert_close(" q", ref_q, tri_q, ratio=1e-5) + assert_close(" k", ref_k, tri_k, ratio=1e-5) + assert_close("dq", ref_dq, tri_dq, ratio=1e-5) + assert_close("dk", ref_dk, tri_dk, ratio=1e-5) + + +@pytest.mark.parametrize("B", [2]) +@pytest.mark.parametrize("T", [2048, 4096]) +@pytest.mark.parametrize("H", [4]) +@pytest.mark.parametrize("G", [1, 4]) +@pytest.mark.parametrize("D", [128, 256]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_rotary_with_offsets(B: int, T: int, H: int, G: int, D: int, dtype: torch.dtype): + torch.manual_seed(42) + q = torch.randn(B, T, H, D).to(device).to(dtype=dtype).requires_grad_() + k = torch.randn(B, T, H//G, D).to(device).to(dtype=dtype).requires_grad_() + seqlen_offset = torch.randint(0, T//2, (B,)).to(device) + max_seqlen = T + seqlen_offset.max().item() + rotary = RotaryEmbedding(D).to(device) + + tri_q, tri_k = rotary(q, k, seqlen_offset=seqlen_offset, max_seqlen=max_seqlen) + tri_dq = torch.autograd.grad(tri_q.sum(), q, retain_graph=True)[0] + tri_dk = torch.autograd.grad(tri_k.sum(), k, retain_graph=True)[0] + + ref_q = torch.cat([ + rotary_embedding_ref( + q[i:i+1].float(), + rotary._cos_cached[offset:offset+T], + rotary._sin_cached[offset:offset+T], + ) + for i, offset in enumerate(seqlen_offset.tolist()) + ]).to(dtype=dtype) + ref_k = torch.cat([ + rotary_embedding_ref( + k[i:i+1].float(), + rotary._cos_cached[offset:offset+T], + rotary._sin_cached[offset:offset+T], + ) + for i, offset in enumerate(seqlen_offset.tolist()) + ]).to(dtype=dtype) + ref_dq = torch.autograd.grad(ref_q.sum(), q, retain_graph=True)[0] + ref_dk = torch.autograd.grad(ref_k.sum(), k, retain_graph=True)[0] + + assert_close(" q", ref_q, tri_q, ratio=1e-5) + assert_close(" k", ref_k, tri_k, ratio=1e-5) + assert_close("dq", ref_dq, tri_dq, ratio=1e-5) + assert_close("dk", ref_dk, tri_dk, ratio=1e-5) + + +@pytest.mark.parametrize("N", [4]) +@pytest.mark.parametrize("T", [2048, 4096]) +@pytest.mark.parametrize("H", [4]) +@pytest.mark.parametrize("G", [1, 4]) +@pytest.mark.parametrize("D", [128, 256]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_rotary_varlen(N: int, T: int, H: int, G: int, D: int, dtype: torch.dtype): + torch.manual_seed(42) + q = torch.randn(1, T, H, D).to(device).to(dtype=dtype).requires_grad_() + k = torch.randn(1, T, H//G, D).to(device).to(dtype=dtype).requires_grad_() + cu_seqlens = torch.cat([ + torch.tensor([0], dtype=torch.long), + torch.arange(1, T)[torch.randperm(T - 1)[:N-1]], + torch.tensor([T], dtype=torch.long), + ], 0).to(device).sort()[0] + rotary = RotaryEmbedding(D).to(device) + + tri_q, tri_k = rotary(q, k, cu_seqlens=cu_seqlens) + tri_dq = torch.autograd.grad(tri_q.sum(), q, retain_graph=True)[0] + tri_dk = torch.autograd.grad(tri_k.sum(), k, retain_graph=True)[0] + + ref_q = torch.cat([ + rotary_embedding_ref( + q[0, start:end].float(), + rotary._cos_cached[:end-start], + rotary._sin_cached[:end-start], + ) + for start, end in zip(cu_seqlens.tolist(), cu_seqlens[1:].tolist(), strict=False) + ]).to(dtype=dtype).unsqueeze(0) + ref_k = torch.cat([ + rotary_embedding_ref( + k[0, start:end].float(), + rotary._cos_cached[:end-start], + rotary._sin_cached[:end-start], + ) + for start, end in zip(cu_seqlens.tolist(), cu_seqlens[1:].tolist(), strict=False) + ]).to(dtype=dtype).unsqueeze(0) + ref_dq = torch.autograd.grad(ref_q.sum(), q, retain_graph=True)[0] + ref_dk = torch.autograd.grad(ref_k.sum(), k, retain_graph=True)[0] + + assert_close(" q", ref_q, tri_q, ratio=1e-5) + assert_close(" k", ref_k, tri_k, ratio=1e-5) + assert_close("dq", ref_dq, tri_dq, ratio=1e-5) + assert_close("dk", ref_dk, tri_dk, ratio=1e-5) diff --git a/code/flash-linear-attention/tests/modules/test_token_shift.py b/code/flash-linear-attention/tests/modules/test_token_shift.py new file mode 100644 index 0000000000000000000000000000000000000000..7729023ede3be83874eb185c85bcc2a167c9cc8a --- /dev/null +++ b/code/flash-linear-attention/tests/modules/test_token_shift.py @@ -0,0 +1,133 @@ + +import pytest +import torch + +from fla.modules.token_shift import token_shift, token_shift_ref +from fla.utils import assert_close, device + +test_b_list = [4] +test_t_list = [512, 4100, 8192] +test_h_list = [2560, 4096] +test_cu_seqlens_list = [ + None, + [0, 4, 7, 40, 128], + [0, 10, 20, 64], + [0, 32], + [0, 1, 3, 4], +] +test_dtype_list = [torch.float] + + +@pytest.mark.parametrize('B', test_b_list) +@pytest.mark.parametrize('T', test_t_list) +@pytest.mark.parametrize('H', test_h_list) +@pytest.mark.parametrize('cu_seqlens_val', test_cu_seqlens_list) +@pytest.mark.parametrize('dtype', test_dtype_list) +def test_token_shift(B, T, H, cu_seqlens_val, dtype): + if cu_seqlens_val is not None: + B = 1 + T = cu_seqlens_val[-1] + cu_seqlens_tensor = torch.tensor(cu_seqlens_val, dtype=torch.int32, device=device) + else: + cu_seqlens_tensor = None + + torch.manual_seed(42) + + x = torch.randn(B, T, H, device=device).to(dtype).requires_grad_(True) + dy = torch.randn_like(x) + + ref = token_shift_ref(x, cu_seqlens_tensor) + tri = token_shift(x, cu_seqlens_tensor) + + ref.backward(dy) + ref_dx, x.grad = x.grad, None + + tri.backward(dy) + tri_dx, x.grad = x.grad, None + + assert_close(' x', ref, tri, 1e-3) + assert_close('dx', ref_dx, tri_dx, 1e-3) + + +def _split_for_passing( + x: torch.Tensor, + cu_seqlens, + split_at: int = 1, +): + assert x.size(0) == 1 + assert 0 < split_at < len(cu_seqlens) - 1 + + cu0 = [t - cu_seqlens[0] for t in cu_seqlens[: split_at + 1]] + cu1 = [t - cu_seqlens[split_at] for t in cu_seqlens[split_at:]] + T0, T1 = cu0[-1], cu1[-1] + + x0 = x[:, :T0].contiguous() + x1 = x[:, T0: T0 + T1].contiguous() + cache1 = x[:, T0 - 1: T0].contiguous() + return x0, x1, \ + torch.tensor(cu0, dtype=torch.int32, device=x.device), \ + torch.tensor(cu1, dtype=torch.int32, device=x.device), \ + cache1 + + +def _check_passing_vs_whole( + B: int, + T: int, + H: int, + cu_seqlens: list[int] | None, + dtype: torch.dtype, + split_at: int = 1, +): + torch.manual_seed(42) + + if cu_seqlens is None: + x = torch.randn(B, T, H, device=device, dtype=dtype, requires_grad=True) + cu_seqlens_tensor = None + else: + B = 1 + T = cu_seqlens[-1] + x = torch.randn(1, T, H, device=device, dtype=dtype, requires_grad=True) + cu_seqlens_tensor = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + + dy = torch.randn_like(x) + ref_out = token_shift(x, cu_seqlens_tensor) + ref_out.backward(dy) + ref_dx = x.grad.clone() + x.grad.zero_() + + if cu_seqlens is None: + T0 = T // 2 + x0 = x[:, :T0].contiguous() + x1 = x[:, T0:].contiguous() + cu0, cu1 = None, None + else: + if split_at >= len(cu_seqlens) - 1: + pytest.skip("invalid split_at") + x0, x1, cu0, cu1, cache1 = _split_for_passing(x, cu_seqlens, split_at) + + out0, cache_out0 = token_shift(x0, cu0, output_cache=True) + out1, cache_out1 = token_shift(x1, cu1, cache=cache_out0, output_cache=True) + + cat_out = torch.cat([out0, out1], dim=1) + cat_out.backward(dy) + + cat_dx = x.grad.clone() + + assert_close("do", ref_out, cat_out, 1e-3) + assert_close("dx", ref_dx, cat_dx, 1e-3) + + +@pytest.mark.parametrize( + ("B", "T", "H", "cu_seqlens", "split_at"), + [ + pytest.param(*test, id="B{}-T{}-H{}-cu{}-split{}".format(*test)) + for test in [ + (2, 512, 1024, None, 1), + (1, 8192, 1024, None, 2), + ] + ], +) +def test_all_with_and_without_varlen(B, T, H, cu_seqlens, split_at): + dtype = torch.float + assert cu_seqlens is None, "This test is for cu_seqlens=None case" + _check_passing_vs_whole(B, T, H, cu_seqlens, dtype, split_at) diff --git a/code/flash-linear-attention/tests/ops/test_attn.py b/code/flash-linear-attention/tests/ops/test_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..58e121837cc5b9d505199d27ce454dc82de3143d --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_attn.py @@ -0,0 +1,125 @@ + +import os + +import pytest +import torch + +from fla.ops.attn.parallel import parallel_attn +from fla.ops.utils import prepare_lens +from fla.utils import assert_close, check_shared_mem, device + +try: + from flash_attn import flash_attn_func, flash_attn_varlen_func + HAS_FLASH = True +except Exception: + HAS_FLASH = False + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'HQ', 'D', 'scale'), + [ + pytest.param(*test, id="B{}-T{}-H{}-HQ{}-D{}-scale{}".format(*test)) + for test in [ + (1, 63, 1, 1, 64, 1.0), + (3, 111, 2, 2, 100, 1.0), + (3, 1024, 2, 8, 60, 0.1), + (3, 1024, 2, 8, 128, 0.1), + (4, 2048, 2, 8, 64, 0.1), + ] + ], +) +def test_parallel( + B: int, + T: int, + H: int, + HQ: int, + D: int, + scale: float, +): + if not check_shared_mem('hopper') and D > 128: + pytest.skip(reason="Skip test, do not have enough shard mem") + if not HAS_FLASH: + pytest.skip(reason="Skipping test because flash-attn is not installed") + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + q = torch.randn((B, T, HQ, D), dtype=torch.float16, device=device).requires_grad_(True) + k = torch.randn((B, T, H, D), dtype=torch.float16, device=device).requires_grad_(True) + v = torch.randn((B, T, H, D), dtype=torch.float16, device=device).requires_grad_(True) + do = torch.randn((B, T, HQ, D), dtype=torch.float16, device=device) + + ref = flash_attn_func(q=q, k=k, v=v, softmax_scale=scale, causal=True) + ref.backward(do) + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + + tri = parallel_attn(q=q, k=k, v=v, scale=scale) + tri.backward(do) + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + + assert_close(" o", ref, tri, 0.005) + assert_close("dq", ref_dq, tri_dq, 0.005) + assert_close("dk", ref_dk, tri_dk, 0.005) + assert_close("dv", ref_dv, tri_dv, 0.005) + + +@pytest.mark.parametrize( + ('H', 'HQ', 'D', 'cu_seqlens'), + [ + pytest.param(*test, id="H{}-HQ{}-D{}-cu_seqlens{}".format(*test)) + for test in [ + (2, 2, 64, [0, 15]), + (2, 8, 64, [0, 256, 500, 1000]), + (2, 2, 100, [0, 15, 100, 300, 1200, 2000]), + ] + ], +) +def test_parallel_varlen( + H: int, + HQ: int, + D: int, + cu_seqlens: list[int], +): + if not HAS_FLASH: + pytest.skip(reason="Skipping test because flash-attn is not installed") + T = cu_seqlens[-1] + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + dtype = torch.float16 + + q = torch.randn((1, T, HQ, D), dtype=dtype, device=device).requires_grad_() + k = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + do = torch.randn((1, T, HQ, D), dtype=dtype, device=device) + + ref = flash_attn_varlen_func( + q=q.squeeze(0), + k=k.squeeze(0), + v=v.squeeze(0), + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=prepare_lens(cu_seqlens).max(), + max_seqlen_k=prepare_lens(cu_seqlens).max(), + causal=True, + ) + ref.backward(do.squeeze(0)) + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + + tri = parallel_attn( + q=q, + k=k, + v=v, + cu_seqlens=cu_seqlens, + ) + tri.backward(do) + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + + assert_close(" o", ref, tri, 0.004) + assert_close("dq", ref_dq.squeeze(), tri_dq.squeeze(), 0.005) + assert_close("dk", ref_dk.squeeze(), tri_dk.squeeze(), 0.005) + assert_close("dv", ref_dv.squeeze(), tri_dv.squeeze(), 0.005) diff --git a/code/flash-linear-attention/tests/ops/test_based.py b/code/flash-linear-attention/tests/ops/test_based.py new file mode 100644 index 0000000000000000000000000000000000000000..7ae4087e0cbd37a1350a46efde575292c1fdd284 --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_based.py @@ -0,0 +1,63 @@ + +import pytest +import torch + +from fla.ops.based import fused_chunk_based, parallel_based +from fla.ops.based.naive import naive_parallel_based +from fla.utils import device + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (1, 63, 1, 60, torch.float16), + (3, 111, 2, 64, torch.float16), + (3, 1024, 4, 100, torch.float16), + (3, 1024, 8, 128, torch.float16), + (4, 2048, 8, 256, torch.float16), + ] + ], +) +def test_based( + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): + torch.manual_seed(42) + q = torch.randn((B, H, T, 16), dtype=dtype, device=device).requires_grad_() + k = torch.randn((B, H, T, 16), dtype=dtype, device=device).requires_grad_() + v = torch.randn((B, H, T, D), dtype=dtype, device=device).requires_grad_() + do = torch.randn_like(v) + ref = naive_parallel_based(q, k, v, use_norm=True) + ref.backward(do) + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + + tri = parallel_based(q, k, v, use_norm=True) + tri.backward(do) + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + + if dtype == torch.float32: + assert ref.allclose(tri, 0, 1e-4) + assert ref_dq.allclose(tri_dq, 0, 1e-4) + assert ref_dk.allclose(tri_dk, 0, 1e-4) + assert ref_dv.allclose(tri_dv, 0, 1e-4) + + tri = fused_chunk_based(q, k, v, use_norm=True) + tri.backward(do) + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + + if dtype == torch.float32: + assert ref.allclose(tri, 0, 1e-4) + assert ref_dq.allclose(tri_dq, 0, 1e-4) + assert ref_dk.allclose(tri_dk, 0, 1e-4) + assert ref_dv.allclose(tri_dv, 0, 1e-4) diff --git a/code/flash-linear-attention/tests/ops/test_comba.py b/code/flash-linear-attention/tests/ops/test_comba.py new file mode 100644 index 0000000000000000000000000000000000000000..4e6b775e4357852919a0cd94c92b2bc598a1cfb2 --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_comba.py @@ -0,0 +1,368 @@ + +import os + +import pytest +import torch +import torch.nn.functional as F +from einops import rearrange + +from fla.ops.comba import chunk_comba, fused_recurrent_comba +from fla.ops.comba.utils import chunk_comba_cumsum_scalar_fwd +from fla.utils import assert_close, device, is_intel_alchemist + + +def cumsum_comba_local_fwd_reference(s, reverse=False, chunk_size=128): + o_0 = torch.zeros_like(s) + o_1 = torch.zeros_like(s) + T = s.size(1) + fn = torch.cumsum + for i in range(0, T, chunk_size): + s_chunk = s[:, i:i+chunk_size] + o_1[:, i:i+chunk_size] = fn(s_chunk.float(), dim=1).to(o_1) + o_0[:, i:i+chunk_size] = o_1[:, i:i+chunk_size] - s_chunk + + return o_0, o_1 + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'chunk_size', 'dtype'), + [ + pytest.param(*test, id='B{}-T{}-H{}-chunk_size{}-{}'.format(*test)) + for test in [ + (32, 200, 4, 64, torch.float), + (32, 1000, 4, 64, torch.float), + (32, 2048, 8, 128, torch.float), + ] + ], +) +def test_cumsum_local_scalar_fwd( + B: int, + T: int, + H: int, + chunk_size: int, + dtype: torch.dtype, +): + s = torch.randn((B, T, H), dtype=dtype, device=device).requires_grad_() + ref_0, ref_1 = cumsum_comba_local_fwd_reference(s, chunk_size=chunk_size) + tri_0, tri_1 = chunk_comba_cumsum_scalar_fwd(s, chunk_size=chunk_size) + assert_close("local cumsum scalar", ref_0, tri_0, 0.001 if dtype == torch.float else 0.003) + assert_close("local cumsum scalar", ref_1, tri_1, 0.001 if dtype == torch.float else 0.003) + + +def chunk_comba_ref( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + p: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + chunk_size: int = 64, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, +): + BT = chunk_size + if scale is None: + scale = 1 / (q.shape[-1] ** 0.5) + # Calculate padding needed to make T a multiple of BT + q, k, v, p, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, p, beta, g]) + + T = q.shape[-2] + pad_len = (BT - (T % BT)) % BT + if pad_len > 0: + # Pad all tensors + q = F.pad(q, (0, 0, 0, pad_len)) + k = F.pad(k, (0, 0, 0, pad_len)) + v = F.pad(v, (0, 0, 0, pad_len)) + p = F.pad(p, (0, 0, 0, pad_len)) + beta = F.pad(beta, (0, pad_len)) + g = F.pad(g, (0, pad_len)) + decay = g + chunk_size = BT + b, h, l, d_k = q.shape + d_v = v.shape[-1] + q = q * scale + v = v * beta[..., None] + p_beta = p * beta[..., None] + assert l % chunk_size == 0 + # note that diagonal is masked. + mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0) + q, k, v, p_beta, decay, g = map( + lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size), + [q, k, v, p_beta, decay.unsqueeze(-1), g.unsqueeze(-1)], + ) + decay = decay.squeeze(-1).cumsum(-1) # [B, H, n, c] + decay_0 = decay - g.squeeze(-1) # [B, H, n, c] + L_mask = ((decay.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().float()).tril() + L_mask_0 = ((decay_0.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().float()).tril() + # [B, H, n, c, d] @ [B, H, n, d, c] -> [B, H, n, c, c] + attn = -((p_beta @ k.transpose(-1, -2)) * L_mask_0).masked_fill(mask, 0) + for i in range(1, chunk_size): + attn[..., i, :i] = attn[..., i, :i].clone() + (attn[..., i, :i, None].clone() * attn[..., :i, :i].clone()).sum(-2) + attn = attn + torch.eye(chunk_size, dtype=torch.float, device=q.device) + # for U + k_cumsum = attn @ v + # for W + k_cumdecay = attn @ (p_beta * decay_0[..., None].exp()) + v = k_cumsum + S = k.new_zeros(b, h, d_k, d_v) + if initial_state is not None: + S += initial_state + o = torch.zeros_like(v) + mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1) + for i in range(0, l // chunk_size): + q_i, k_i, v_i = q[:, :, i], k[:, :, i], v[:, :, i] + attn = (q_i @ k_i.transpose(-1, -2) * L_mask[:, :, i]).masked_fill_(mask, 0) + v_prime = k_cumdecay[:, :, i] @ S + v_new = v_i - v_prime + o_inter = (q_i * decay[:, :, i, :, None].exp()) @ S + o[:, :, i] = o_inter + attn @ v_new + S = S * decay[:, :, i, -1, None, None].exp() + (k_i * (decay[:, :, i, -1, None] - decay[:, :, i]).exp() + [..., None]).transpose(-1, -2) @ v_new + if not output_final_state: + S = None + # unpad + o = rearrange(o, 'b h n c d -> b h (n c) d') + o = o[:, :, :T] + o = o.transpose(1, 2) + return o, S + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-{}".format(*test)) + for test in [ + (1, 63, 1, 64, 1, 1, torch.float), + (2, 1024, 4, 60, 1, 1, torch.float), + (2, 1024, 8, 128, 1, 0.1, torch.float), + (2, 1024, 8, 128, 0.1, 1, torch.float), + (2, 1024, 8, 128, 1, 10, torch.float), + (4, 2048, 8, 64, 0.1, 1, torch.float), + (2, 1024, 8, 128, 1, 0.1, torch.float16), + (2, 1024, 8, 128, 1, 10, torch.float16), + ] + ], +) +def test_fused_recurrent( + B: int, + T: int, + H: int, + D: int, + scale: float, + gate_logit_normalizer: float, + dtype: torch.dtype, +): + torch.manual_seed(42) + q = F.normalize(torch.randn(B, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype) + k = F.normalize(torch.randn(B, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype) + v = torch.randn(B, T, H, D, dtype=dtype) + p = F.normalize(torch.randn(B, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype) + beta = torch.rand(B, T, H, dtype=dtype).sigmoid() + g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.float32)) + g = g / gate_logit_normalizer + h0 = torch.randn(B, H, D, D, dtype=torch.float32) + q, k, v, p, beta, g, h0 = map(lambda x: x.to(device).requires_grad_(), (q, k, v, p, beta, g, h0)) + ref, ref_ht = chunk_comba_ref( + q=q.clone(), + k=k.clone(), + v=v.clone(), + p=p.clone(), + beta=beta.clone(), + g=g.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + tri, tri_ht = fused_recurrent_comba( + q=q.clone(), + k=k.clone(), + v=v.clone(), + p=p.clone(), + beta=beta.clone(), + g=g.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + assert_close('o', ref, tri, 0.002) + assert_close('ht', ref_ht, tri_ht, 0.002) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'mask_p', 'use_qk_l2norm_in_kernel', 'dtype'), + [ + pytest.param( + *test, + id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-mask_p{}-use_qk_l2norm_in_kernel{}-{}".format(*test), + ) + for test in [ + (1, 63, 1, 64, 1, 1, 0, False, torch.float16), + (2, 1000, 3, 60, 1, 1, 0, False, torch.float16), + (2, 1024, 3, 64, 0.1, 1, 0.5, False, torch.float16), + (2, 1024, 4, 100, 1, 0.1, 0, False, torch.float16), + (2, 1024, 4, 128, 0.1, 1, 0, True, torch.float16), + (2, 1024, 4, 128, 0.1, 1, 0.5, False, torch.float16), + (2, 1024, 4, 128, 0.1, 10, 0, False, torch.float16), + (4, 2048, 8, 64, 0.1, 1, 0, True, torch.float16), + ] + ], +) +def test_chunk( + B: int, + T: int, + H: int, + D: int, + scale: float, + gate_logit_normalizer: float, + mask_p: float, + use_qk_l2norm_in_kernel: bool, + dtype: torch.dtype, +): + if is_intel_alchemist and D > 128: + pytest.skip(reason='chunk_gated_delta_rule is not supported on alchemist for D>128') + + q = torch.randn(B, T, H, D, dtype=dtype) + k = torch.randn(B, T, H, D, dtype=dtype) + p = torch.randn(B, T, H, D, dtype=dtype) + v = torch.randn(B, T, H, D, dtype=dtype) + beta = torch.rand(B, T, H, dtype=dtype).sigmoid() + g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.float32)) + g = g / gate_logit_normalizer + g = g * (torch.rand_like(g) > mask_p) + h0 = torch.zeros(B, H, D, D, dtype=torch.float32) + q, k, v, p, beta, g, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, p, beta, g, h0)) + + tri, tri_ht = chunk_comba( + q=F.normalize(q.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q.clone(), + k=F.normalize(k.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k.clone(), + p=F.normalize(p.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else p.clone(), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) + do = torch.randn_like(v) + dht = torch.randn_like(h0) + ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True) + tri_dq, tri_dk, tri_dv, tri_dp, tri_dbeta, tri_dg, tri_dh0 = q.grad, k.grad, v.grad, p.grad, beta.grad, g.grad, h0.grad + q.grad = k.grad = v.grad = p.grad = beta.grad = g.grad = h0.grad = None + + ref, ref_ht = chunk_comba_ref( + q=F.normalize(q.clone(), p=2, dim=-1), + k=F.normalize(k.clone(), p=2, dim=-1), + p=F.normalize(p.clone(), p=2, dim=-1), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + + ((ref * do).sum() + (ref_ht * dht).sum()).backward() + ref_dq, ref_dk, ref_dv, ref_dp, ref_dbeta, ref_dg, ref_dh0 = q.grad, k.grad, v.grad, p.grad, beta.grad, g.grad, h0.grad + + assert_close(" o", ref, tri, 0.005) + assert_close(" ht", ref_ht, tri_ht, 0.005) + assert_close(" dq", ref_dq, tri_dq, 0.005) + assert_close(" dk", ref_dk, tri_dk, 0.008) + assert_close(" dv", ref_dv, tri_dv, 0.005) + assert_close(" dp", ref_dp, tri_dp, 0.008) + assert_close(" dg", ref_dg, tri_dg, 0.02) + assert_close(" db", ref_dbeta, tri_dbeta, 0.005) + assert_close("dh0", ref_dh0, tri_dh0, 0.008) + + +@pytest.mark.parametrize( + ('H', 'D', 'mask_p', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-D{}-mask_p{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (4, 64, 0, [0, 15], torch.float16), + (4, 64, 0, [0, 256, 500, 1000], torch.float16), + (4, 64, 0.5, [0, 256, 500, 1000], torch.float16), + (4, 100, 0, [0, 15, 100, 300, 1200, 2000], torch.float16), + ] + ], +) +@pytest.mark.skipif( + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set', +) +def test_chunk_varlen( + H: int, + D: int, + mask_p: float, + cu_seqlens: list[int], + dtype: torch.dtype, +): + if is_intel_alchemist and D > 128: + pytest.skip(reason='chunk_gated_delta_rule is not supported on alchemist for D>128') + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + + N = len(cu_seqlens) - 1 + T = cu_seqlens[-1] + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + + q = torch.randn((1, T, H, D), dtype=dtype) + k = F.normalize(torch.randn(1, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype) + v = torch.randn((1, T, H, D), dtype=dtype) + p = F.normalize(torch.randn(1, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype) + g = F.logsigmoid(torch.rand(1, T, H, dtype=dtype)) + g = g * (torch.rand_like(g) > mask_p) + beta = torch.rand(1, T, H, dtype=dtype).sigmoid() + h0 = torch.randn((N, H, D, D), dtype=dtype) + + q, k, v, p, beta, g, h0 = map(lambda x: x.to(device).requires_grad_(), (q, k, v, p, beta, g, h0)) + do = torch.randn_like(v) + dht = torch.rand_like(h0) + + tri, tri_ht = chunk_comba( + q=q.clone(), + k=k.clone(), + v=v.clone(), + p=p.clone(), + beta=beta.clone(), + g=g.clone(), + output_final_state=True, + initial_state=h0.clone(), + cu_seqlens=cu_seqlens, + ) + ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True) + tri_dq, tri_dk, tri_dv, tri_dbeta, tri_dg, tri_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad + q.grad = k.grad = v.grad = beta.grad = g.grad = h0.grad = None + + ref = [] + ref_ht = [] + for i in range(N): + ref_i, ref_ht_i = chunk_comba_ref( + q=q[:, cu_seqlens[i]:cu_seqlens[i+1]], + k=k[:, cu_seqlens[i]:cu_seqlens[i+1]], + v=v[:, cu_seqlens[i]:cu_seqlens[i+1]], + p=p[:, cu_seqlens[i]:cu_seqlens[i+1]], + beta=beta[:, cu_seqlens[i]:cu_seqlens[i+1]], + g=g[:, cu_seqlens[i]:cu_seqlens[i+1]], + initial_state=h0[i], + output_final_state=True, + ) + ref.append(ref_i) + ref_ht.append(ref_ht_i) + ref = torch.cat(ref, 1) + ref_ht = torch.cat(ref_ht, 0) + + ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) + ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dg, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad + + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.007) + assert_close('dk', ref_dk, tri_dk, 0.008) + assert_close('dv', ref_dv, tri_dv, 0.007) + assert_close('db', ref_dbeta, tri_dbeta, 0.015) + assert_close('dg', ref_dg, tri_dg, 0.015) + assert_close('dh0', ref_dh0, tri_dh0, 0.007) diff --git a/code/flash-linear-attention/tests/ops/test_delta.py b/code/flash-linear-attention/tests/ops/test_delta.py new file mode 100644 index 0000000000000000000000000000000000000000..2032208252c9b2fe02fa36e8f737a1009ea8fdfa --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_delta.py @@ -0,0 +1,152 @@ + + +import pytest +import torch +import torch.nn.functional as F + +from fla.ops.delta_rule import chunk_delta_rule, fused_recurrent_delta_rule +from fla.utils import assert_close, device, device_platform + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'scale', 'use_qk_l2norm_in_kernel', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-{}".format(*test)) + for test in [ + (1, 63, 1, 64, 1, False, torch.float16), + (2, 100, 4, 60, 0.1, False, torch.float16), + (2, 1000, 3, 128, 0.1, False, torch.float16), + (2, 1024, 4, 128, 1, True, torch.float16), + (3, 2000, 4, 128, 0.1, False, torch.float16), + (4, 2048, 8, 64, 0.1, False, torch.float16), + ] + ], +) +@pytest.mark.skipif( + device_platform == 'intel', + reason='Intel Triton Failure', +) +def test_chunk( + B: int, + T: int, + H: int, + D: int, + scale: float, + use_qk_l2norm_in_kernel: bool, + dtype: torch.dtype, +): + torch.manual_seed(42) + q = torch.randn(B, T, H, D, dtype=dtype) + k = torch.randn(B, T, H, D, dtype=dtype) + v = torch.randn(B, T, H, D, dtype=dtype) + beta = torch.randn(B, T, H, dtype=dtype).sigmoid() + h0 = torch.randn(B, H, D, D, dtype=torch.float32) + q, k, v, beta, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, beta, h0)) + do = torch.rand_like(v) + dht = torch.rand_like(h0) + + tri, tri_ht = chunk_delta_rule( + q=F.normalize(q.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q.clone(), + k=F.normalize(k.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k.clone(), + v=v.clone(), + beta=beta.clone(), + scale=scale, + output_final_state=True, + initial_state=h0.clone(), + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) + ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True) + tri_dq, tri_dk, tri_dv, tri_dbeta, tri_dh0 = q.grad, k.grad, v.grad, beta.grad, h0.grad + q.grad = k.grad = v.grad = beta.grad = h0.grad = None + + ref, ref_ht = fused_recurrent_delta_rule( + q=F.normalize(q.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q.clone(), + k=F.normalize(k.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k.clone(), + v=v.clone(), + beta=beta.clone(), + scale=scale, + output_final_state=True, + initial_state=h0.clone(), + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) + ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) + ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, h0.grad + + assert_close('o', ref, tri, 0.006) + assert_close('ht', ref_ht, tri_ht, 0.006) + assert_close('dq', ref_dq, tri_dq, 0.008) + assert_close('dk', ref_dk, tri_dk, 0.008) + assert_close('dv', ref_dv, tri_dv, 0.008) + assert_close('db', ref_dbeta, tri_dbeta, 0.008) + assert_close('dh0', ref_dh0, tri_dh0, 0.008) + + +@pytest.mark.parametrize( + ('H', 'D', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-D{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (2, 64, [0, 15], torch.float16), + (3, 60, [0, 111, 500], torch.float16), + (3, 64, [0, 256, 500, 900, 1000], torch.float16), + (4, 100, [0, 15, 100, 300, 1200, 1599, 1800, 2000], torch.float16), + ] + ], +) +@pytest.mark.skipif( + device_platform == 'intel', + reason='Intel Triton Failure', +) +def test_chunk_varlen( + H: int, + D: int, + cu_seqlens: list[int], + dtype: torch.dtype, +): + torch.manual_seed(42) + T = cu_seqlens[-1] + N = len(cu_seqlens) - 1 + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + + # seq-first required for inputs with variable lengths + q = torch.randn((1, T, H, D), dtype=dtype) + k = F.normalize(torch.randn(1, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype) + v = torch.randn((1, T, H, D), dtype=dtype) + beta = torch.randn(1, T, H, dtype=dtype).sigmoid() + h0 = torch.randn(N, H, D, D, dtype=dtype) + q, k, v, beta, h0 = map(lambda x: x.to(device).requires_grad_(), (q, k, v, beta, h0)) + do = torch.randn_like(v) + dht = torch.rand_like(h0) + + ref, ref_ht = fused_recurrent_delta_rule( + q=q.clone(), + k=k.clone(), + v=v.clone(), + beta=beta.clone(), + output_final_state=True, + initial_state=h0.clone(), + cu_seqlens=cu_seqlens, + ) + ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) + ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, h0.grad + + tri, tri_ht = chunk_delta_rule( + q=q.clone(), + k=k.clone(), + v=v.clone(), + beta=beta.clone(), + output_final_state=True, + initial_state=h0.clone(), + cu_seqlens=cu_seqlens, + ) + ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True) + tri_dq, tri_dk, tri_dv, tri_dbeta, tri_dh0 = q.grad, k.grad, v.grad, beta.grad, h0.grad + q.grad = k.grad = v.grad = beta.grad = h0.grad = None + + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.008) + assert_close('dk', ref_dk, tri_dk, 0.008) + assert_close('dv', ref_dv, tri_dv, 0.008) + assert_close('db', ref_dbeta, tri_dbeta, 0.008) + assert_close('dh0', ref_dh0, tri_dh0, 0.008) diff --git a/code/flash-linear-attention/tests/ops/test_delta_product.py b/code/flash-linear-attention/tests/ops/test_delta_product.py new file mode 100644 index 0000000000000000000000000000000000000000..7b5437eb2a817d58c7b98e2bd1d91b04168f66de --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_delta_product.py @@ -0,0 +1,186 @@ + + +import pytest +import torch +import torch.nn.functional as F + +from fla.ops.gated_delta_product import chunk_gated_delta_product +from fla.ops.gated_delta_product.chunk_ref import chunk_gated_delta_product_ref +from fla.ops.gated_delta_product.naive import naive_recurrent_gated_delta_product +from fla.utils import assert_close, device + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'scale', 'num_householder', 'use_qk_l2norm_in_kernel', 'dtype'), + [ + pytest.param( + *test, + id="B{}-T{}-H{}-D{}-scale{}-num_householder{}-l2norm{}-{}".format(*test), + ) + for test in [ + (1, 63, 1, 64, 0.1, 1, False, torch.float16), + (2, 200, 3, 60, 0.1, 1, False, torch.float16), + (2, 1000, 4, 64, 0.1, 2, False, torch.float16), + (2, 1024, 4, 64, 1, 2, True, torch.float16), + (2, 1024, 6, 100, 1, 2, False, torch.float16), + (4, 1500, 8, 128, 0.1, 3, False, torch.float16), + (2, 2048, 8, 128, 1, 3, False, torch.float16), + (2, 2048, 8, 128, 1, 3, True, torch.float16), + ] + ], +) +def test_chunk( + B: int, + T: int, + H: int, + D: int, + scale: float, + num_householder: int, + use_qk_l2norm_in_kernel: bool, + dtype: torch.dtype, +): + torch.manual_seed(42) + q = torch.randn(B, T, H, D, dtype=dtype) + k = torch.randn(B, T * num_householder, H, D, dtype=dtype) + v = torch.randn(B, T * num_householder, H, D, dtype=dtype) + beta = torch.rand(B, T * num_householder, H, dtype=dtype).sigmoid() + h0 = torch.zeros(B, H, D, D, dtype=torch.float32) + q, k, v, beta, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, beta, h0)) + + tri, tri_ht = chunk_gated_delta_product( + q=F.normalize(q.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q.clone(), + k=F.normalize(k.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k.clone(), + v=v.clone(), + g=None, + beta=beta.clone(), + num_householder=num_householder, + scale=scale, + output_final_state=True, + initial_state=h0.clone(), + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) + do = torch.randn_like(q) + dht = torch.randn_like(h0) + ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True) + tri_dq, tri_dk, tri_dv, tri_dbeta, tri_dh0 = q.grad, k.grad, v.grad, beta.grad, h0.grad + q.grad = k.grad = v.grad = beta.grad = h0.grad = None + + ref, ref_ht = chunk_gated_delta_product_ref( + q=F.normalize(q.clone(), p=2, dim=-1), + k=F.normalize(k.clone(), p=2, dim=-1), + v=v.clone(), + g=None, + beta=beta.clone(), + num_householder=num_householder, + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + + ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) + ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, h0.grad + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.008) + assert_close('dk', ref_dk, tri_dk, 0.008) + assert_close('dv', ref_dv, tri_dv, 0.008) + assert_close('db', ref_dbeta, tri_dbeta, 0.02) + assert_close('dh0', ref_dh0, tri_dh0, 0.008) + + +@pytest.mark.parametrize( + ('H', 'D', 'num_householder', 'cu_seqlens', 'dtype'), + [ + (2, 64, 3, [0, 63 ], torch.float16), + (2, 100, 2, [0, 63, 100, 500, 1000], torch.float16), + (2, 128, 2, [0, 100, 300, 800, 1500, 2000], torch.float16), + (2, 256, 3, [0, 100, 123, 300, 500, 800, 1000, 1500, 2048], torch.float16), + ], +) +def test_chunk_varlen( + H: int, + D: int, + num_householder: int, + cu_seqlens: list[int], + dtype: torch.dtype, +): + torch.manual_seed(42) + + T = cu_seqlens[-1] + N = len(cu_seqlens) - 1 + cu_seqlens = torch.LongTensor(cu_seqlens).to(device) + scale = 1.0 + + q = torch.nn.functional.normalize(torch.randn((1, T, H, D), dtype=dtype), dim=-1, p=2) + k = torch.nn.functional.normalize(torch.randn(1, T*num_householder, H, D, dtype=dtype), dim=-1, p=2) + v = torch.randn((1, T*num_householder, H, D), dtype=dtype) + beta = torch.rand(1, T*num_householder, H, dtype=dtype).sigmoid() + h0 = torch.randn((N, H, D, D), dtype=dtype) + + q, k, v, beta, h0 = map(lambda x: x.to(device).requires_grad_(), (q, k, v, beta, h0)) + do = torch.randn_like(q) + dht = torch.rand_like(h0) + + tri, tri_ht = chunk_gated_delta_product( + q=q.clone(), + k=k.clone(), + v=v.clone(), + beta=beta.clone(), + g=None, + scale=scale, + output_final_state=True, + num_householder=num_householder, + initial_state=h0.clone(), + cu_seqlens=cu_seqlens, + ) + ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True) + tri_dq, tri_dk, tri_dv, tri_dbeta, tri_dh0 = q.grad, k.grad, v.grad, beta.grad, h0.grad + q.grad = k.grad = v.grad = beta.grad = h0.grad = None + + ref, ref_ht = chunk_gated_delta_product_ref( + q=q.clone(), + k=k.clone(), + v=v.clone(), + beta=beta.clone(), + g=None, + scale=scale, + output_final_state=True, + num_householder=num_householder, + initial_state=h0.clone(), + cu_seqlens=cu_seqlens, + ) + ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) + ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, h0.grad + + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.007) + assert_close('dk', ref_dk, tri_dk, 0.008) + assert_close('dv', ref_dv, tri_dv, 0.007) + assert_close('db', ref_dbeta, tri_dbeta, 0.015) + assert_close('dh0', ref_dh0, tri_dh0, 0.007) + q.grad = k.grad = v.grad = beta.grad = h0.grad = None + + torch_ref = torch.zeros_like(ref) + torch_ref_ht = torch.zeros_like(ref_ht) + for i in range(len(cu_seqlens) - 1): + start, end = cu_seqlens[i], cu_seqlens[i+1] + q_i = q[:, start:end, :, :] + k_i = k[:, start*num_householder:end*num_householder, :, :] + v_i = v[:, start*num_householder:end*num_householder, :, :] + beta_i = beta[:, start*num_householder:end*num_householder, :] + o3_i, h3_i = naive_recurrent_gated_delta_product( + q_i, k_i, v_i, None, beta_i, scale=scale, cu_seqlens=None, output_final_state=True, num_householder=num_householder, + ) + torch_ref[:, start:end, :, :] = o3_i + torch_ref_ht[i, :, :, :] = h3_i.squeeze(0) + + ((torch_ref * do).sum() + (torch_ref_ht * dht).sum()).backward(retain_graph=True) + + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.007) + assert_close('dk', ref_dk, tri_dk, 0.008) + assert_close('dv', ref_dv, tri_dv, 0.007) + assert_close('db', ref_dbeta, tri_dbeta, 0.015) + assert_close('dh0', ref_dh0, tri_dh0, 0.007) diff --git a/code/flash-linear-attention/tests/ops/test_deltaformer.py b/code/flash-linear-attention/tests/ops/test_deltaformer.py new file mode 100644 index 0000000000000000000000000000000000000000..7ea0f5918eae7d91dac564827eea01dcc4df3358 --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_deltaformer.py @@ -0,0 +1,130 @@ + + +import pytest +import torch + +from fla.ops.deltaformer import deltaformer_attn +from fla.ops.deltaformer.naive import naive_deltaformer_attn +from fla.utils import assert_close, device, is_intel_alchemist + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 128, 2, 64, torch.float16), + (1, 256, 4, 64, torch.float16), + (2, 512, 4, 64, torch.float16), + (4, 1024, 4, 128, torch.float16), + ] + ], +) +@pytest.mark.skipif( + is_intel_alchemist, + reason="Skipping test on Intel Alchemist due to known issues with SRAM.", +) +def test_deltaformer_attn( + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): + """ + Test DeltaFormer pre-attention by comparing fused implementation with naive reference. + """ + torch.manual_seed(42) + + q = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True) + k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True) + v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True) + beta = torch.randn((B, T, H), dtype=dtype, device=device).sigmoid().requires_grad_(True) + + do = torch.randn((B, T, H, D), dtype=dtype, device=device) + + ref = naive_deltaformer_attn(q, k, v, beta) + ref.backward(do) + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_dbeta, beta.grad = beta.grad.clone(), None + + tri = deltaformer_attn(q, k, v, beta) + tri.backward(do) + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_dbeta, beta.grad = beta.grad.clone(), None + + assert_close('o', ref, tri, 0.006) + assert_close('dq', ref_dq, tri_dq, 0.008) + assert_close('dk', ref_dk, tri_dk, 0.008) + assert_close('dv', ref_dv, tri_dv, 0.008) + assert_close('dbeta', ref_dbeta, tri_dbeta, 0.008) + + +@pytest.mark.parametrize( + ('H', 'D', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-D{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (2, 64, [0, 63], torch.float16), + (4, 64, [0, 256, 500, 1000], torch.float16), + (4, 128, [0, 15, 100, 300, 1200, 2000], torch.float16), + (2, 128, [0, 100, 123, 300, 500, 800, 1000, 1500, 2048], torch.float16), + ] + ], +) +@pytest.mark.skipif( + is_intel_alchemist, + reason="Skipping test on Intel Alchemist due to known issues with SRAM.", +) +def test_deltaformer_attn_varlen( + H: int, + D: int, + cu_seqlens: list[int], + dtype: torch.dtype, +): + torch.manual_seed(42) + + T = cu_seqlens[-1] + N = len(cu_seqlens) - 1 + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + + q = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + k = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + beta = torch.randn((1, T, H), dtype=dtype, device=device).sigmoid().requires_grad_() + + do = torch.randn_like(q) + + refs = [] + for i in range(N): + ref = naive_deltaformer_attn( + q[:, cu_seqlens[i]:cu_seqlens[i+1]], + k[:, cu_seqlens[i]:cu_seqlens[i+1]], + v[:, cu_seqlens[i]:cu_seqlens[i+1]], + beta[:, cu_seqlens[i]:cu_seqlens[i+1]], + ) + refs.append(ref) + ref = torch.cat(refs, dim=1) + + ref.backward(do) + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_dbeta, beta.grad = beta.grad.clone(), None + + tri = deltaformer_attn(q, k, v, beta, cu_seqlens=cu_seqlens) + tri.backward(do) + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_dbeta, beta.grad = beta.grad.clone(), None + + assert_close('o', ref, tri, 0.006) + assert_close('dq', ref_dq, tri_dq, 0.008) + assert_close('dk', ref_dk, tri_dk, 0.008) + assert_close('dv', ref_dv, tri_dv, 0.008) + assert_close('dbeta', ref_dbeta, tri_dbeta, 0.008) diff --git a/code/flash-linear-attention/tests/ops/test_dplr_delta.py b/code/flash-linear-attention/tests/ops/test_dplr_delta.py new file mode 100644 index 0000000000000000000000000000000000000000..1239b4a70281256f071b0147c23b58ef094f4200 --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_dplr_delta.py @@ -0,0 +1,432 @@ + +import os + +import pytest +import torch +import torch.nn.functional as F +from einops import rearrange + +from fla.ops.generalized_delta_rule.dplr import chunk_dplr_delta_rule, fused_recurrent_dplr_delta_rule +from fla.utils import assert_close, device, device_platform + + +def recurrent_dplr_delta_rule_ref( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + gk: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, +): + q, k, v, a, b, gk = map(lambda x: x.transpose(1, 2).to(torch.float), (q, k, v, a, b, gk)) + + B, H, T, K, V = *q.shape, v.shape[-1] + o = torch.zeros_like(v) + S = torch.zeros(B, H, K, V).to(v) + if initial_state is not None: + S = initial_state + if scale is None: + scale = K ** -0.5 + q = q * scale + + for i in range(T): + _q = q[:, :, i] + _k = k[:, :, i] + _v = v[:, :, i].clone() + a_i = a[:, :, i] + b_i = b[:, :, i] + # first matmul then decay in DPLR. + _v2 = (S.clone() * a_i[..., None]).sum(-2) + S = S.clone() * gk[:, :, i].exp()[..., None] + S = S.clone() + _k.unsqueeze(-1) * _v.unsqueeze(-2) + b_i.unsqueeze(-1) * _v2.unsqueeze(-2) + o[:, :, i] = torch.einsum('bhd,bhdm->bhm', _q, S) + if not output_final_state: + S = None + o = o.transpose(1, 2) + return o, S + + +def chunk_dplr_delta_rule_ref( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + gk: torch.Tensor, + initial_state: torch.Tensor = None, + output_final_state: bool = True, + scale: float = None, + chunk_size: int = 64, +): + q, k, v, a, b, gk = map(lambda x: x.transpose(1, 2).to(torch.float), (q, k, v, a, b, gk)) + BT = chunk_size + T = q.shape[-2] + pad_len = (BT - (T % BT)) % BT + + q, k, v, a, b, gk = map(lambda x: F.pad(x, (0, 0, 0, pad_len)).to(torch.float), [q, k, v, a, b, gk]) + B, H, _, K, V = *q.shape, v.shape[-1] + NT = q.shape[-2] // BT + if scale is None: + scale = K ** -0.5 + q = q * scale + + S = k.new_zeros(B, H, K, V) + if initial_state is not None: + S += initial_state + + # note that diagonal is masked. + mask = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=0) + q, k, v, a, b, gk = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=BT), [q, k, v, a, b, gk]) + gk_cumsum = gk.cumsum(-2) + A_ab = torch.zeros(B, H, NT, BT, BT).to(q.device) + A_qk = torch.zeros(B, H, NT, BT, BT).to(q.device) + A_ak = torch.zeros(B, H, NT, BT, BT).to(q.device) + A_qb = torch.zeros(B, H, NT, BT, BT).to(q.device) + + for i in range(BT): + a_i = a[:, :, :, i, None] + q_i = q[:, :, :, i, None] + gk_i = gk_cumsum[:, :, :, i, None] + mask = (torch.arange(BT) <= i).to(q.device) + attn_i = (gk_i - gk_cumsum).masked_fill(~mask.unsqueeze(-1), float('-inf')).exp() + A_qk[:, :, :, i, :] = (q_i * k * attn_i).sum(-1).clone() + A_qb[:, :, :, i, :] = (q_i * b * attn_i).sum(-1).clone() + mask = (torch.arange(BT) < i).to(q.device) + # shift by one. + attn_i = (gk_i - gk[:, :, :, i, None] - gk_cumsum).masked_fill(~mask.unsqueeze(-1), float('-inf')).exp() + A_ab[:, :, :, i, :] = (a_i * b * attn_i).sum(-1).clone() + A_ak[:, :, :, i, :] = (a_i * k * attn_i).sum(-1).clone() + + A_ab = A_ab + for i in range(1, BT): + A_ab[..., i, :i] = A_ab[..., i, :i].clone() + (A_ab[..., i, :, None].clone() * A_ab[..., :, :i].clone()).sum(-2) + + A_ab = A_ab + torch.eye(BT, dtype=torch.float, device=q.device) + u = A_ab @ (A_ak @ v) + w = A_ab @ ((gk_cumsum-gk).exp() * a) + + o = torch.zeros_like(v) + mask = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=1) + for i in range(0, NT): + q_i, k_i, v_i, u_i, w_i, b_i = q[:, :, i], k[:, :, i], v[:, :, i], u[:, :, i], w[:, :, i], b[:, :, i] + v2_i = u_i + w_i @ S + o_1 = A_qk[:, :, i] @ v_i + o_2 = A_qb[:, :, i] @ v2_i + o_3 = (q_i * gk_cumsum[:, :, i].exp()) @ S + o[:, :, i] = o_1 + o_2 + o_3 + decay = (gk_cumsum[:, :, i, -1, None] - gk_cumsum[:, :, i]).exp() + S = S*gk_cumsum[:, :, i, -1, :, None].exp() + (k_i * decay).transpose(-1, -2) @ v_i + \ + (b_i * decay).transpose(-1, -2) @ v2_i + + S = None if output_final_state is False else S + o = rearrange(o, 'b h n c d -> b h (n c) d') + o = o[:, :, :T].transpose(1, 2) + return o, S + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'scale', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-{}".format(*test)) + for test in [ + (1, 63, 1, 64, 1, torch.float), + (2, 1024, 4, 60, 1, torch.float), + (2, 1024, 8, 128, 1, torch.float), + (2, 1024, 8, 128, 0.1, torch.float), + (4, 2048, 8, 64, 0.1, torch.float), + (2, 1024, 8, 128, 1, torch.float16), + ] + ], +) +def test_recurrent_fwd( + B: int, + T: int, + H: int, + D: int, + scale: float, + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + os.environ['TORCH_CUDA_MATMUL_PRECISION'] = 'highest' + q = torch.randn(B, T, H, D, dtype=dtype) + k = torch.randn(B, T, H, D, dtype=dtype) + v = torch.randn(B, T, H, D, dtype=dtype) + a = torch.rand(B, T, H, D, dtype=dtype) + gk = torch.randn(B, T, H, D, dtype=torch.float) + + a = F.normalize(a, p=2, dim=-1) + b = -a + gk = F.logsigmoid(gk) / 16 + + h0 = torch.randn(B, H, D, D, dtype=torch.float) + q, k, v, a, b, gk, h0 = map(lambda x: x.to(device).requires_grad_(False), (q, k, v, a, b, gk, h0)) + ref, ref_ht = chunk_dplr_delta_rule_ref( + q=q.clone(), + k=k.clone(), + v=v.clone(), + a=a.clone(), + b=b.clone(), + gk=gk.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + tri, tri_ht = recurrent_dplr_delta_rule_ref( + q=q.clone(), + k=k.clone(), + v=v.clone(), + a=a.clone(), + b=b.clone(), + gk=gk.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + assert_close('o', ref, tri, 0.001) + assert_close('ht', ref_ht, tri_ht, 0.001) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'scale', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-{}".format(*test)) + for test in [ + (1, 63, 1, 64, 1, torch.float), + (2, 1024, 4, 60, 1, torch.float), + (2, 1024, 8, 100, 1, torch.float), + (2, 1024, 8, 128, 0.1, torch.float), + (4, 2048, 8, 64, 0.1, torch.float), + ] + ], +) +def test_fused_recurrent( + B: int, + T: int, + H: int, + D: int, + scale: float, + dtype: torch.dtype, +): + torch.manual_seed(42) + q = torch.randn(B, T, H, D, dtype=dtype) + k = torch.randn(B, T, H, D, dtype=dtype) + v = torch.randn(B, T, H, D, dtype=dtype) + a = torch.rand(B, T, H, D, dtype=dtype) + gk = torch.randn(B, T, H, D, dtype=torch.float) + + a = F.normalize(a, p=2, dim=-1) + b = -a + gk = F.logsigmoid(gk) / 4 + + h0 = torch.randn(B, H, D, D, dtype=torch.float) + q, k, v, a, b, gk, h0 = map(lambda x: x.to(device).requires_grad_(False), (q, k, v, a, b, gk, h0)) + ref, ref_ht = recurrent_dplr_delta_rule_ref( + q=q.clone(), + k=k.clone(), + v=v.clone(), + a=a.clone(), + b=b.clone(), + gk=gk.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + + tri, tri_ht = fused_recurrent_dplr_delta_rule( + q=q.clone(), + k=k.clone(), + v=v.clone(), + a=a.clone(), + b=b.clone(), + gk=gk.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + assert_close('o', ref, tri, 0.002) + assert_close('ht', ref_ht, tri_ht, 0.002) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'mask_p', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-mask_p{}-{}".format(*test)) + for test in [ + (1, 63, 1, 64, 1, 1, 0, torch.float16), + (2, 1000, 3, 60, 1, 1, 0, torch.float16), + (2, 1024, 3, 64, 0.1, 1, 0.5, torch.float16), + (2, 1024, 4, 100, 1, 0.1, 0, torch.float16), + (2, 1024, 4, 128, 0.1, 1, 0, torch.float16), + (2, 1024, 4, 128, 0.1, 1, 0.5, torch.float16), + (2, 1024, 4, 128, 0.1, 10, 0, torch.float16), + (4, 2048, 8, 64, 0.1, 1, 0, torch.float16), + ] + ], +) +@pytest.mark.skipif( + device_platform == 'intel', + reason='Intel Triton Failure', +) +def test_chunk( + B: int, + T: int, + H: int, + D: int, + scale: float, + gate_logit_normalizer: float, + mask_p: float, + dtype: torch.dtype, +): + torch.manual_seed(42) + q = torch.randn(B, T, H, D, dtype=dtype) + k = torch.randn(B, T, H, D, dtype=dtype) + v = torch.randn(B, T, H, D, dtype=dtype) + a = torch.rand(B, T, H, D, dtype=dtype) + gk = torch.randn(B, T, H, D, dtype=torch.float) + + a = F.normalize(a, p=2, dim=-1) + b = -a + gk = F.logsigmoid(gk) + gk = gk / gate_logit_normalizer + gk = gk * (torch.rand_like(gk) > mask_p) + + h0 = torch.randn(B, H, D, D, dtype=torch.float) + q, k, v, a, b, gk, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, a, b, gk, h0)) + ref, ref_ht = chunk_dplr_delta_rule_ref( + q=q.clone(), + k=k.clone(), + v=v.clone(), + a=a.clone(), + b=b.clone(), + gk=gk.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + do = torch.randn_like(v) + dht = torch.randn_like(h0) + ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) + ref_dq, ref_dk, ref_dv, ref_da, ref_db, ref_dg, ref_dh0 = q.grad, k.grad, v.grad, a.grad, b.grad, gk.grad, h0.grad + q.grad = k.grad = v.grad = a.grad = b.grad = gk.grad = h0.grad = None + + tri, tri_ht = chunk_dplr_delta_rule( + q=q.clone(), + k=k.clone(), + v=v.clone(), + a=a.clone(), + b=b.clone(), + gk=gk.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True) + tri_dq, tri_dk, tri_dv, tri_da, tri_db, tri_dg, tri_dh0 = q.grad, k.grad, v.grad, a.grad, b.grad, gk.grad, h0.grad + q.grad = k.grad = v.grad = a.grad = b.grad = gk.grad = h0.grad = None + + assert_close('o', ref, tri, 0.007) + assert_close('ht', ref_ht, tri_ht, 0.008) + assert_close('dq', ref_dq, tri_dq, 0.008) + assert_close('dk', ref_dk, tri_dk, 0.008) + assert_close('dv', ref_dv, tri_dv, 0.008) + assert_close('da', ref_da, tri_da, 0.008) + assert_close('db', ref_db, tri_db, 0.008) + if gate_logit_normalizer >= 1 and ref_dg.norm() > 0.01: # otherwise it is meaningless + assert_close('dg', ref_dg, tri_dg, 0.008) + assert_close('dh0', ref_dh0, tri_dh0, 0.008) + + +@pytest.mark.parametrize( + ('H', 'D', 'mask_p', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-D{}-mask_p{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (4, 64, 0, [0, 15], torch.float16), + (4, 64, 0, [0, 256, 500, 1000], torch.float16), + (4, 64, 0.5, [0, 256, 500, 1000], torch.float16), + (4, 100, 0, [0, 15, 100, 300, 1111, 1599, 2000], torch.float16), + ] + ], +) +@pytest.mark.skipif( + device_platform == 'intel', + reason='Intel Triton Failure', +) +def test_chunk_varlen( + H: int, + D: int, + mask_p: float, + cu_seqlens: list[int], + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + + N = len(cu_seqlens) - 1 + T = cu_seqlens[-1] + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + + # seq-first required for inputs with variable lengths + q = torch.randn(1, T, H, D, dtype=dtype) + k = torch.randn(1, T, H, D, dtype=dtype) + v = torch.randn(1, T, H, D, dtype=dtype) + a = torch.rand(1, T, H, D, dtype=dtype) + gk = torch.randn(1, T, H, D, dtype=torch.float) + a = F.normalize(a, p=2, dim=-1) + b = -a + gk = F.logsigmoid(gk) + gk = gk * (torch.rand_like(gk) > mask_p) + h0 = torch.randn(N, H, D, D, dtype=torch.float) + q, k, v, a, b, gk, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, a, b, gk, h0)) + + tri, tri_ht = chunk_dplr_delta_rule( + q=q.clone(), + k=k.clone(), + v=v.clone(), + a=a.clone(), + b=b.clone(), + gk=gk.clone(), + output_final_state=True, + initial_state=h0.clone(), + cu_seqlens=cu_seqlens, + ) + do = torch.randn_like(v) + dht = torch.randn_like(h0) + ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True) + tri_dq, tri_dk, tri_dv, tri_da, tri_db, tri_dg, tri_dh0 = q.grad, k.grad, v.grad, a.grad, b.grad, gk.grad, h0.grad + q.grad = k.grad = v.grad = a.grad = b.grad = gk.grad = h0.grad = None + + ref = [] + ref_ht = [] + for i in range(N): + ref_i, ref_ht_i = chunk_dplr_delta_rule_ref( + q=q[:, cu_seqlens[i]:cu_seqlens[i+1]], + k=k[:, cu_seqlens[i]:cu_seqlens[i+1]], + v=v[:, cu_seqlens[i]:cu_seqlens[i+1]], + a=a[:, cu_seqlens[i]:cu_seqlens[i+1]], + b=b[:, cu_seqlens[i]:cu_seqlens[i+1]], + gk=gk[:, cu_seqlens[i]:cu_seqlens[i+1]], + initial_state=h0[i, None], + output_final_state=True, + ) + ref.append(ref_i) + ref_ht.append(ref_ht_i) + + ref = torch.cat(ref, 1) + ref_ht = torch.cat(ref_ht, 0) + ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) + ref_dq, ref_dk, ref_dv, ref_da, ref_db, ref_dg, ref_dh0 = q.grad, k.grad, v.grad, a.grad, b.grad, gk.grad, h0.grad + + assert_close('o', ref, tri, 0.007) + assert_close('ht', ref_ht, tri_ht, 0.008) + assert_close('dq', ref_dq, tri_dq, 0.008) + assert_close('dk', ref_dk, tri_dk, 0.008) + assert_close('dv', ref_dv, tri_dv, 0.008) + assert_close('da', ref_da, tri_da, 0.008) + assert_close('db', ref_db, tri_db, 0.008) + assert_close('dg', ref_dg, tri_dg, 0.008) + assert_close('dh0', ref_dh0, tri_dh0, 0.008) diff --git a/code/flash-linear-attention/tests/ops/test_forgetting_attn.py b/code/flash-linear-attention/tests/ops/test_forgetting_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..f8932b3adeb041cf48c5036e6a520eefb8fca624 --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_forgetting_attn.py @@ -0,0 +1,151 @@ + + +import pytest +import torch +import torch.nn.functional as F +from einops import rearrange, repeat + +from fla.ops.forgetting_attn.parallel import parallel_forgetting_attn +from fla.utils import assert_close, check_shared_mem, device, is_intel_alchemist + + +def naive_forgetting_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + scale: float | None = None, +): + _, T, HQ, D = q.shape + H = k.shape[2] + G = HQ // H + if scale is None: + scale = D ** -0.5 + gc = g.float().cumsum(1) + mask = torch.tril(torch.ones((T, T), dtype=torch.bool, device=device)) + ref = torch.einsum("bqhd,bkhd->bhqk", q.float() * scale, repeat(k, "b t h d -> b t (h g) d", g=G).float()) + ref = ref + rearrange(gc, "b t h -> b h t 1") - rearrange(gc, "b t h -> b h 1 t") + ref = ref.masked_fill(~mask.unsqueeze(0).unsqueeze(0), -float('inf')) + ref = torch.einsum("bhqk,bkhd->bqhd", F.softmax(ref, dim=-1), repeat(v, "b t h d -> b t (h g) d", g=G).float()) + return ref + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'HQ', 'D', 'scale'), + [ + pytest.param(*test, id="B{}-T{}-H{}-HQ{}-D{}-scale{}".format(*test)) + for test in [ + (1, 63, 1, 1, 64, 1.0), + (3, 111, 2, 2, 100, 1.0), + (3, 1024, 2, 8, 60, 0.1), + (3, 1024, 2, 8, 128, 0.1), + (4, 2048, 2, 8, 64, 0.1), + ] + ], +) +def test_parallel( + B: int, + T: int, + H: int, + HQ: int, + D: int, + scale: float, +): + torch.manual_seed(42) + dtype = torch.float16 + if not check_shared_mem('hopper') and D > 128: + # maybe we can enable this test on Triton 3.3.0 + pytest.skip("Skipping test because global shared memory is not available") + + q = torch.randn((B, T, HQ, D), dtype=dtype, device=device).requires_grad_(True) + k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True) + v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True) + + g = torch.randn((B, T, HQ), dtype=dtype, device=device).uniform_(-0.1, -0.01).requires_grad_(True) + + do = torch.randn((B, T, HQ, D), dtype=dtype, device=device) + ref = naive_forgetting_attn(q, k, v, g, scale) + ref.backward(do) + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_dg, g.grad = g.grad.clone(), None + + tri = parallel_forgetting_attn(q=q, k=k, v=v, g=g, scale=scale) + tri.backward(do) + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_dg, g.grad = g.grad.clone(), None + + assert_close(" o", ref, tri, 0.005) + assert_close("dq", ref_dq, tri_dq, 0.005) + assert_close("dk", ref_dk, tri_dk, 0.005) + assert_close("dv", ref_dv, tri_dv, 0.005) + assert_close("dg", ref_dg, tri_dg, 0.005) + + +@pytest.mark.parametrize( + ('H', 'HQ', 'D', 'cu_seqlens'), + [ + pytest.param(*test, id="H{}-HQ{}-D{}-cu_seqlens{}".format(*test)) + for test in [ + (2, 2, 64, [0, 15]), + (2, 8, 64, [0, 256, 500, 1000]), + (2, 2, 100, [0, 15, 100, 300, 1200, 2000]), + ] + ], +) +@pytest.mark.skipif( + is_intel_alchemist, + reason="Intel Triton Failure", +) +def test_parallel_varlen( + H: int, + HQ: int, + D: int, + cu_seqlens: list[int], +): + torch.manual_seed(42) + T = cu_seqlens[-1] + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + dtype = torch.float16 + # seq-first required for inputs with variable lengths + q = torch.randn((1, T, HQ, D), dtype=dtype, device=device).requires_grad_() + k = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + g = torch.rand((1, T, HQ), dtype=dtype, device=device).uniform_(-0.1, -0.01).requires_grad_(True) + do = torch.randn((1, T, HQ, D), dtype=dtype, device=device) + + ref = q.new_empty(1, T, HQ, D) + for bos, eos in zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False): + ref[:, bos:eos] = naive_forgetting_attn( + q=q[:, bos:eos], + k=k[:, bos:eos], + v=v[:, bos:eos], + g=g[:, bos:eos], + ) + ref.backward(do) + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_dg, g.grad = g.grad.clone(), None + + tri = parallel_forgetting_attn( + q=q, + k=k, + v=v, + g=g, + cu_seqlens=cu_seqlens, + ) + tri.backward(do) + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_dg, g.grad = g.grad.clone(), None + + assert_close(" o", ref, tri, 0.004) + assert_close(" dq", ref_dq.squeeze(), tri_dq.squeeze(), 0.005) + assert_close(" dk", ref_dk.squeeze(), tri_dk.squeeze(), 0.005) + assert_close(" dv", ref_dv.squeeze(), tri_dv.squeeze(), 0.005) + assert_close(" dg", ref_dg.squeeze(), tri_dg.squeeze(), 0.005) diff --git a/code/flash-linear-attention/tests/ops/test_gated_delta.py b/code/flash-linear-attention/tests/ops/test_gated_delta.py new file mode 100644 index 0000000000000000000000000000000000000000..96406c50031fb082d0fc6098361b0b1299f29be7 --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_gated_delta.py @@ -0,0 +1,353 @@ + +import os + +import pytest +import torch +import torch.nn.functional as F +from einops import rearrange, repeat + +from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule +from fla.utils import assert_close, device, is_intel_alchemist + + +def recurrent_gated_delta_rule_ref( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + g: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, +): + q, k, v, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g]) + B, H, T, K, V = *k.shape, v.shape[-1] + o = torch.zeros(B, H, T, V).to(v) + h = torch.zeros(B, H, K, V).to(v) + if initial_state is not None: + h = initial_state + if scale is None: + scale = 1 / (q.shape[-1] ** 0.5) + q = q * scale + for i in range(T): + b_q = q[:, :, i] + b_k = k[:, :, i] + b_v = v[:, :, i].clone() + h = h.clone() * g[:, :, i].exp()[..., None, None] + b_beta = beta[:, :, i] + b_v = b_v - (h.clone() * b_k[..., None]).sum(-2) + b_v = b_v * b_beta[..., None] + h = h.clone() + b_k.unsqueeze(-1) * b_v.unsqueeze(-2) + o[:, :, i] = torch.einsum('bhd,bhdm->bhm', b_q, h) + if not output_final_state: + h = None + o = o.transpose(1, 2).contiguous() + return o, h + + +def chunk_gated_delta_rule_ref( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + chunk_size: int = 64, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, +): + BT = chunk_size + if scale is None: + scale = 1 / (q.shape[-1] ** 0.5) + # Calculate padding needed to make T a multiple of BT + q, k, v, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g]) + + T = q.shape[-2] + pad_len = (BT - (T % BT)) % BT + if pad_len > 0: + # Pad all tensors + q = F.pad(q, (0, 0, 0, pad_len)) + k = F.pad(k, (0, 0, 0, pad_len)) + v = F.pad(v, (0, 0, 0, pad_len)) + beta = F.pad(beta, (0, pad_len)) + g = F.pad(g, (0, pad_len)) + q, k, v, beta, g = map(lambda x: x.to(torch.float32), [q, k, v, beta, g]) + decay = g + chunk_size = BT + b, h, l, d_k = q.shape + d_v = v.shape[-1] + q = q * scale + v = v * beta[..., None] + k_beta = k * beta[..., None] + assert l % chunk_size == 0 + # note that diagonal is masked. + mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0) + q, k, v, k_beta, decay = map( + lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size), + [q, k, v, k_beta, decay.unsqueeze(-1)], + ) + decay = decay.squeeze(-1).cumsum(-1) + decay_exp = decay.exp()[..., None] + L_mask = ((decay.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().float()).tril() + attn = -((k_beta @ k.transpose(-1, -2)) * L_mask).masked_fill(mask, 0) + for i in range(1, chunk_size): + attn[..., i, :i] = attn[..., i, :i].clone() + (attn[..., i, :i, None].clone() * attn[..., :i, :i].clone()).sum(-2) + attn = attn + torch.eye(chunk_size, dtype=torch.float, device=q.device) + attn = attn + k_cumsum = attn @ v + k_cumdecay = attn @ (k_beta * decay_exp) + v = k_cumsum + S = k.new_zeros(b, h, d_k, d_v) + if initial_state is not None: + S = initial_state + o = torch.zeros_like(v) + mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1) + for i in range(0, l // chunk_size): + q_i, k_i, v_i = q[:, :, i], k[:, :, i], v[:, :, i] + attn = (q_i @ k_i.transpose(-1, -2) * L_mask[:, :, i]).masked_fill_(mask, 0) + v_prime = (k_cumdecay[:, :, i]) @ S + v_new = v_i - v_prime + o_inter = (q_i * decay[:, :, i, :, None].exp()) @ S + o[:, :, i] = o_inter + attn @ v_new + S = S * decay[:, :, i, -1, None, None].exp() + (k_i * (decay[:, :, i, -1, None] - decay[:, :, i]).exp() + [..., None]).transpose(-1, -2) @ v_new + if not output_final_state: + S = None + # unpad + o = rearrange(o, 'b h n c d -> b h (n c) d') + o = o[:, :, :T] + o = o.transpose(1, 2) + return o, S + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'HV', 'D', 'scale', 'gate_logit_normalizer', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-HV{}-D{}-scale{}-gate_logit_normalizer{}-{}".format(*test)) + for test in [ + (1, 63, 1, 1, 64, 1, 1, torch.float), + (2, 500, 4, 4, 60, 1, 1, torch.float), + (2, 1000, 2, 8, 128, 1, 0.1, torch.float), + (3, 1024, 2, 2, 128, 0.1, 1, torch.float), + (4, 1024, 3, 3, 128, 1, 10, torch.float), + (4, 2048, 4, 4, 64, 0.1, 1, torch.float), + (2, 1024, 4, 4, 128, 1, 0.1, torch.float16), + (2, 1024, 4, 8, 128, 1, 10, torch.float16), + ] + ], +) +def test_fused_recurrent( + B: int, + T: int, + H: int, + HV: int, + D: int, + scale: float, + gate_logit_normalizer: float, + dtype: torch.dtype, +): + torch.manual_seed(42) + q = torch.randn(B, T, H, D, dtype=torch.float32) + k = torch.randn(B, T, H, D, dtype=torch.float32) + v = torch.randn(B, T, HV, D, dtype=dtype) + beta = torch.rand(B, T, HV, dtype=dtype).sigmoid() + g = F.logsigmoid(torch.rand(B, T, HV, dtype=torch.float32)) + g = g / gate_logit_normalizer + h0 = torch.randn(B, HV, D, D, dtype=torch.float32) + q, k, v, beta, g, h0 = map(lambda x: x.to(device).requires_grad_(), (q, k, v, beta, g, h0)) + ref, ref_ht = recurrent_gated_delta_rule_ref( + q=F.normalize(repeat(q.clone(), 'b t h d -> b t (h g) d', g=HV // H), p=2, dim=-1).to(dtype), + k=F.normalize(repeat(k.clone(), 'b t h d -> b t (h g) d', g=HV // H), p=2, dim=-1).to(dtype), + v=v.clone(), + beta=beta.clone(), + g=g.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + tri, tri_ht = fused_recurrent_gated_delta_rule( + q=q.clone(), + k=k.clone(), + v=v.clone(), + beta=beta.clone(), + g=g.clone(), + scale=scale, + initial_state=h0.clone(), + use_qk_l2norm_in_kernel=True, + output_final_state=True, + ) + assert_close('o', ref, tri, 0.002) + assert_close('ht', ref_ht, tri_ht, 0.002) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'mask_p', 'use_qk_l2norm_in_kernel', 'dtype'), + [ + pytest.param( + *test, + id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-mask_p{}-use_qk_l2norm_in_kernel{}-{}".format(*test), + ) + for test in [ + (1, 63, 1, 64, 1, 1, 0, False, torch.float16), + (2, 500, 3, 60, 1, 1, 0, False, torch.float16), + (2, 1000, 3, 64, 0.1, 1, 0.5, False, torch.float16), + (3, 1024, 4, 100, 1, 0.1, 0, False, torch.float16), + (4, 1024, 4, 128, 0.1, 1, 0, False, torch.float16), + (4, 1024, 4, 128, 0.1, 1, 0, True, torch.float16), + (2, 1500, 4, 128, 0.1, 10, 0, False, torch.float16), + (4, 2048, 8, 64, 0.1, 1, 0, False, torch.float16), + ] + ], +) +def test_chunk( + B: int, + T: int, + H: int, + D: int, + scale: float, + gate_logit_normalizer: float, + mask_p: float, + use_qk_l2norm_in_kernel: bool, + dtype: torch.dtype, +): + torch.manual_seed(42) + if is_intel_alchemist and D > 128: + pytest.skip(reason='chunk_gated_delta_rule is not supported on alchemist for D>128') + + q = torch.rand(B, T, H, D, dtype=dtype) + k = torch.rand(B, T, H, D, dtype=dtype) + v = torch.rand(B, T, H, D, dtype=dtype) + beta = torch.rand(B, T, H, dtype=dtype).sigmoid() + g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.float32)) + g = g / gate_logit_normalizer + g = g * (torch.rand_like(g) > mask_p) + h0 = torch.zeros(B, H, D, D, dtype=torch.float32) + q, k, v, beta, g, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, beta, g, h0)) + + tri, tri_ht = chunk_gated_delta_rule( + q=F.normalize(q.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q.clone(), + k=F.normalize(k.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k.clone(), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) + do = torch.randn_like(v) + dht = torch.randn_like(h0) + ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True) + tri_dq, tri_dk, tri_dv, tri_dbeta, tri_dg, tri_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad + q.grad = k.grad = v.grad = beta.grad = g.grad = h0.grad = None + + ref, ref_ht = recurrent_gated_delta_rule_ref( + q=F.normalize(q.clone(), p=2, dim=-1), + k=F.normalize(k.clone(), p=2, dim=-1), + v=v.clone(), + beta=beta.clone(), + g=g.clone(), + scale=scale, + output_final_state=True, + initial_state=h0.clone(), + ) + + ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) + ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dg, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.008) + assert_close('dk', ref_dk, tri_dk, 0.008) + assert_close('dv', ref_dv, tri_dv, 0.008) + assert_close('db', ref_dbeta, tri_dbeta, 0.02) + assert_close('dg', ref_dg, tri_dg, 0.02) + assert_close('dh0', ref_dh0, tri_dh0, 0.008) + + +@pytest.mark.parametrize( + ('H', 'D', 'mask_p', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-D{}-mask_p{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (4, 60, 0, [0, 15], torch.float16), + (4, 64, 0, [0, 256, 500, 1000], torch.float16), + (4, 64, 0.5, [0, 256, 500, 1000], torch.float16), + (4, 100, 0, [0, 15, 100, 300, 1200, 2000], torch.float16), + ] + ], +) +@pytest.mark.skipif( + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set', +) +def test_chunk_varlen( + H: int, + D: int, + mask_p: float, + cu_seqlens: list[int], + dtype: torch.dtype, +): + if is_intel_alchemist and D > 128: + pytest.skip(reason='chunk_gated_delta_rule is not supported on alchemist for D>128') + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + # randomly split the sequence into N segments + cu_seqlens = torch.LongTensor(cu_seqlens).to(device) + T = cu_seqlens[-1] + N = len(cu_seqlens) - 1 + + # seq-first required for inputs with variable lengths + q = torch.randn((1, T, H, D), dtype=dtype) + k = F.normalize(torch.randn(1, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype) + v = torch.randn((1, T, H, D), dtype=dtype) + g = F.logsigmoid(torch.rand(1, T, H, dtype=dtype)) + g = g * (torch.rand_like(g) > mask_p) + beta = torch.rand(1, T, H, dtype=dtype).sigmoid() + h0 = torch.randn((N, H, D, D), dtype=dtype) + + q, k, v, beta, g, h0 = map(lambda x: x.to(device).requires_grad_(), (q, k, v, beta, g, h0)) + do = torch.randn_like(v) + dht = torch.rand_like(h0) + + tri, tri_ht = chunk_gated_delta_rule( + q=q.clone(), + k=k.clone(), + v=v.clone(), + beta=beta.clone(), + g=g.clone(), + initial_state=h0.clone(), + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True) + tri_dq, tri_dk, tri_dv, tri_dbeta, tri_dg, tri_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad + q.grad = k.grad = v.grad = beta.grad = g.grad = h0.grad = None + + ref = [] + ref_ht = [] + for i in range(N): + ref_i, ref_ht_i = recurrent_gated_delta_rule_ref( + q=q[:, cu_seqlens[i]:cu_seqlens[i+1]], + k=k[:, cu_seqlens[i]:cu_seqlens[i+1]], + v=v[:, cu_seqlens[i]:cu_seqlens[i+1]], + beta=beta[:, cu_seqlens[i]:cu_seqlens[i+1]], + g=g[:, cu_seqlens[i]:cu_seqlens[i+1]], + initial_state=h0[i], + output_final_state=True, + ) + ref.append(ref_i) + ref_ht.append(ref_ht_i) + ref = torch.cat(ref, 1) + ref_ht = torch.cat(ref_ht, 0) + + ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) + ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dg, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad + + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.007) + assert_close('dk', ref_dk, tri_dk, 0.008) + assert_close('dv', ref_dv, tri_dv, 0.007) + assert_close('db', ref_dbeta, tri_dbeta, 0.015) + assert_close('dg', ref_dg, tri_dg, 0.015) + assert_close('dh0', ref_dh0, tri_dh0, 0.007) diff --git a/code/flash-linear-attention/tests/ops/test_gated_delta_product.py b/code/flash-linear-attention/tests/ops/test_gated_delta_product.py new file mode 100644 index 0000000000000000000000000000000000000000..852285b9915b532279ed2a53967939291b821b10 --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_gated_delta_product.py @@ -0,0 +1,210 @@ + +import os + +import pytest +import torch +import torch.nn.functional as F + +from fla.ops.gated_delta_product import chunk_gated_delta_product +from fla.ops.gated_delta_product.chunk_ref import chunk_gated_delta_product_ref +from fla.ops.gated_delta_product.naive import naive_recurrent_gated_delta_product +from fla.utils import assert_close, device, is_intel_alchemist + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'scale', 'num_householder', 'gate_logit_normalizer', 'mask_p', 'use_qk_l2norm_in_kernel', 'dtype'), + [ + pytest.param( + *test, + id="B{}-T{}-H{}-D{}-scale{}-num_householder{}-gate_logit_normalizer{}-mask_p{}-l2norm{}-{}".format(*test), + ) + for test in [ + (1, 63, 1, 64, 0.1, 1, 1, 0, False, torch.float16), + (2, 200, 3, 60, 0.1, 1, 1, 0, False, torch.float16), + (2, 1000, 4, 64, 0.1, 2, 0.1, 0.5, False, torch.float16), + (2, 1024, 4, 64, 1, 2, 1, 0, True, torch.float16), + (2, 1024, 6, 100, 1, 2, 10, 0, False, torch.float16), + (4, 1500, 8, 128, 0.1, 3, 1, 0.5, False, torch.float16), + (2, 2048, 8, 128, 1, 3, 1, 0, False, torch.float16), + (2, 2048, 8, 128, 1, 3, 1, 0, True, torch.float16), + ] + ], +) +def test_chunk( + B: int, + T: int, + H: int, + D: int, + scale: float, + num_householder: int, + gate_logit_normalizer: float, + mask_p: float, + use_qk_l2norm_in_kernel: bool, + dtype: torch.dtype, +): + if is_intel_alchemist and D > 128: + pytest.skip(reason='chunk_gated_delta_rule is not supported on alchemist for D>128') + + q = torch.randn(B, T, H, D, dtype=dtype) + k = torch.randn(B, T * num_householder, H, D, dtype=dtype) + v = torch.randn(B, T * num_householder, H, D, dtype=dtype) + beta = torch.rand(B, T * num_householder, H, dtype=dtype).sigmoid() + g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.float32)) + h0 = torch.zeros(B, H, D, D, dtype=torch.float32) + g = g / gate_logit_normalizer + g = g * (torch.rand_like(g) > mask_p) + q, k, v, beta, g, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, beta, g, h0)) + + tri, tri_ht = chunk_gated_delta_product( + q=F.normalize(q.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q.clone(), + k=F.normalize(k.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k.clone(), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + num_householder=num_householder, + scale=scale, + output_final_state=True, + initial_state=h0.clone(), + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) + do = torch.randn_like(q) + dht = torch.randn_like(h0) + ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True) + tri_dq, tri_dk, tri_dv, tri_dbeta, tri_dg, tri_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad + q.grad = k.grad = v.grad = beta.grad = g.grad = h0.grad = None + + ref, ref_ht = chunk_gated_delta_product_ref( + q=F.normalize(q.clone(), p=2, dim=-1), + k=F.normalize(k.clone(), p=2, dim=-1), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + num_householder=num_householder, + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + + ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) + ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dg, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.008) + assert_close('dk', ref_dk, tri_dk, 0.008) + assert_close('dv', ref_dv, tri_dv, 0.008) + assert_close('db', ref_dbeta, tri_dbeta, 0.02) + assert_close('dg', ref_dg, tri_dg, 0.02) + assert_close('dh0', ref_dh0, tri_dh0, 0.008) + + +@pytest.mark.parametrize( + ('H', 'D', 'num_householder', 'mask_p', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-D{}-num_householder{}-mask_p{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (2, 64, 3, 0, [0, 63], torch.float16), + (2, 100, 2, 0, [0, 63, 100, 500, 1000], torch.float16), + (2, 100, 2, 0, [0, 100, 256, 512, 1500, 1500], torch.float16), + (2, 128, 2, 0, [0, 100, 300, 800, 1500, 2000], torch.float16), + (2, 128, 2, 0.5, [0, 31, 111, 799, 1000, 1500, 1800, 2000], torch.float16), + (2, 128, 2, 0.5, [0, 63, 300, 800, 1000, 1399, 2048], torch.float16), + (2, 256, 3, 0, [0, 100, 123, 300, 500, 800, 1000, 1500, 2048], torch.float16), + ] + ], +) +def test_chunk_varlen( + H: int, + D: int, + num_householder: int, + mask_p: float, + cu_seqlens: list[int], + dtype: torch.dtype, +): + if is_intel_alchemist and D > 128: + pytest.skip(reason='chunk_gated_delta_rule is not supported on alchemist for D>128') + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + cu_seqlens = torch.LongTensor(cu_seqlens).to(device) + T = cu_seqlens[-1] + N = len(cu_seqlens) - 1 + + q = torch.nn.functional.normalize(torch.randn((1, T, H, D), dtype=dtype), dim=-1, p=2) + k = torch.nn.functional.normalize(torch.randn(1, T*num_householder, H, D, dtype=dtype), dim=-1, p=2) + v = torch.randn((1, T*num_householder, H, D), dtype=dtype) + g = F.logsigmoid(torch.rand(1, T, H, dtype=dtype)) + g = g * (torch.rand_like(g) > mask_p) + beta = torch.rand(1, T*num_householder, H, dtype=dtype).sigmoid() + h0 = torch.randn((N, H, D, D), dtype=dtype) + + q, k, v, beta, g, h0 = map(lambda x: x.to(device).requires_grad_(), (q, k, v, beta, g, h0)) + do = torch.randn_like(q) + dht = torch.rand_like(h0) + scale = D ** -0.5 + + tri, tri_ht = chunk_gated_delta_product( + q=q.clone(), + k=k.clone(), + v=v.clone(), + beta=beta.clone(), + g=g.clone(), + scale=scale, + output_final_state=True, + num_householder=num_householder, + initial_state=h0.clone(), + cu_seqlens=cu_seqlens, + ) + ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True) + tri_dq, tri_dk, tri_dv, tri_dbeta, tri_dg, tri_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad + q.grad = k.grad = v.grad = beta.grad = g.grad = h0.grad = None + + ref, ref_ht = chunk_gated_delta_product_ref( + q=q.clone(), + k=k.clone(), + v=v.clone(), + beta=beta.clone(), + g=g.clone(), + scale=scale, + output_final_state=True, + num_householder=num_householder, + initial_state=h0.clone(), + cu_seqlens=cu_seqlens, + ) + + ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) + ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dg, ref_dh0 = q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad + + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.007) + assert_close('dk', ref_dk, tri_dk, 0.008) + assert_close('dv', ref_dv, tri_dv, 0.007) + assert_close('db', ref_dbeta, tri_dbeta, 0.015) + assert_close('dh0', ref_dh0, tri_dh0, 0.007) + assert_close('dg', ref_dg, tri_dg, 0.015) + q.grad = k.grad = v.grad = beta.grad = g.grad = h0.grad = None + + torch_ref = torch.zeros_like(ref) + torch_ref_ht = torch.zeros_like(ref_ht) + for i in range(len(cu_seqlens) - 1): + start, end = cu_seqlens[i], cu_seqlens[i+1] + q_i = q[:, start:end, :, :] + k_i = k[:, start*num_householder:end*num_householder, :, :] + v_i = v[:, start*num_householder:end*num_householder, :, :] + g_i = g[:, start:end, :] + beta_i = beta[:, start*num_householder:end*num_householder, :] + o3_i, h3_i = naive_recurrent_gated_delta_product( + q_i, k_i, v_i, g_i, beta_i, scale=scale, cu_seqlens=None, output_final_state=True, num_householder=num_householder, + ) + torch_ref[:, start:end, :, :] = o3_i + torch_ref_ht[i, :, :, :] = h3_i.squeeze(0) + + ((torch_ref * do).sum() + (torch_ref_ht * dht).sum()).backward(retain_graph=True) + + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.007) + assert_close('dk', ref_dk, tri_dk, 0.008) + assert_close('dv', ref_dv, tri_dv, 0.007) + assert_close('db', ref_dbeta, tri_dbeta, 0.015) + assert_close('dg', ref_dg, tri_dg, 0.015) + assert_close('dh0', ref_dh0, tri_dh0, 0.007) diff --git a/code/flash-linear-attention/tests/ops/test_gla.py b/code/flash-linear-attention/tests/ops/test_gla.py new file mode 100644 index 0000000000000000000000000000000000000000..8306a88896a79a1f67b7ec1c2b34a117b0208d62 --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_gla.py @@ -0,0 +1,325 @@ + +import os + +import pytest +import torch +import torch.nn.functional as F + +from fla.ops.gla import chunk_gla, fused_recurrent_gla +from fla.ops.gla.naive import naive_recurrent_gla +from fla.utils import assert_close, device, device_platform + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'gate_logit_normalizer', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-gate_logit_normalizer{}-{}".format(*test)) + for test in [ + (1, 63, 1, 64, 1, torch.float), + (2, 1024, 4, 60, 1, torch.float), + (2, 1024, 8, 128, 0.1, torch.float), + (2, 1024, 8, 128, 1, torch.float), + (2, 1024, 8, 128, 10, torch.float), + (4, 2048, 8, 64, 1, torch.float), + (2, 1024, 8, 128, 0.1, torch.float16), + (2, 1024, 8, 128, 10, torch.float16), + ] + ], +) +@pytest.mark.skipif( + device_platform == 'intel', + reason='Intel Triton Failure', +) +def test_fused_recurrent( + B: int, + T: int, + H: int, + D: int, + gate_logit_normalizer: float, + dtype: torch.dtype, +): + torch.manual_seed(42) + + q = torch.rand((B, T, H, D), dtype=dtype, device=device).requires_grad_() + k = torch.rand((B, T, H, D), dtype=dtype, device=device).requires_grad_() + v = torch.rand((B, T, H, D), dtype=dtype, device=device).requires_grad_() + g = (F.logsigmoid(torch.rand((B, T, H, D), dtype=dtype, device=device)) / gate_logit_normalizer).requires_grad_() + h0 = torch.rand(B, H, D, D, device=device).requires_grad_() + do = torch.randn_like(v) + dht = torch.randn((B, H, D, D), dtype=dtype, device=device) + + ref, ref_ht = naive_recurrent_gla( + q=q, + k=k, + v=v, + gk=g, + initial_state=h0, + output_final_state=True, + ) + ((ref * do).sum() + (ref_ht * dht).sum()).backward() + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_dg, g.grad = g.grad.clone(), None + ref_dh0, h0.grad = h0.grad.clone(), None + + tri, tri_ht = fused_recurrent_gla( + q=q, + k=k, + v=v, + gk=g, + initial_state=h0, + output_final_state=True, + ) + ((tri * do).sum() + (tri_ht * dht).sum()).backward() + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_dg, g.grad = g.grad.clone(), None + tri_dh0, h0.grad = h0.grad.clone(), None + + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) + assert_close('dg', ref_dg, tri_dg, 0.005) + assert_close('dh0', ref_dh0, tri_dh0, 0.005) + + +@pytest.mark.parametrize( + ('H', 'D', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-D{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (4, 64, [0, 15], torch.float), + (4, 64, [0, 256, 500, 1000], torch.float), + (4, 100, [0, 15, 100, 300, 1200, 2000], torch.float), + (4, 64, [0, 1, 100, 300, 1200, 2048], torch.float16), + (4, 128, [0, 200, 512, 1200, 2048], torch.float16), + ] + ], +) +@pytest.mark.skipif( + device_platform == 'intel', + reason='Intel Triton Failure', +) +def test_fused_recurrent_varlen( + H: int, + D: int, + cu_seqlens: list[int], + dtype: torch.dtype, +): + torch.manual_seed(42) + + N = len(cu_seqlens) - 1 + T = cu_seqlens[-1] + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + + q = torch.rand((1, T, H, D), dtype=dtype, device=device).requires_grad_() + k = torch.rand((1, T, H, D), dtype=dtype, device=device).requires_grad_() + v = torch.rand((1, T, H, D), dtype=dtype, device=device).requires_grad_() + g = F.logsigmoid(torch.rand((1, T, H, D), dtype=dtype, device=device)).requires_grad_() + h0 = torch.rand(N, H, D, D, device=device).requires_grad_() + do = torch.randn_like(v) + dht = torch.randn((N, H, D, D), dtype=dtype, device=device) + + refs, ref_hts = [], [] + for i in range(N): + ref, ref_ht = naive_recurrent_gla( + q=q[:, cu_seqlens[i]:cu_seqlens[i+1]], + k=k[:, cu_seqlens[i]:cu_seqlens[i+1]], + v=v[:, cu_seqlens[i]:cu_seqlens[i+1]], + gk=g[:, cu_seqlens[i]:cu_seqlens[i+1]], + initial_state=h0[i], + output_final_state=True, + ) + refs.append(ref) + ref_hts.append(ref_ht) + ref = torch.cat(refs, dim=1) + ref_ht = torch.cat(ref_hts, dim=0) + + ((ref * do).sum() + (ref_ht * dht).sum()).backward() + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_dg, g.grad = g.grad.clone(), None + ref_dh0, h0.grad = h0.grad.clone(), None + + tri, tri_ht = fused_recurrent_gla( + q=q, + k=k, + v=v, + gk=g, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + ((tri * do).sum() + (tri_ht * dht).sum()).backward() + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_dg, g.grad = g.grad.clone(), None + tri_dh0, h0.grad = h0.grad.clone(), None + + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) + assert_close('dg', ref_dg, tri_dg, 0.005) + assert_close('dh0', ref_dh0, tri_dh0, 0.005) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'gate_logit_normalizer', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-gate_logit_normalizer{}-{}".format(*test)) + for test in [ + (1, 63, 1, 64, 1, torch.float16), + (2, 1024, 4, 60, 1, torch.float16), + (2, 1024, 8, 128, 0.1, torch.float16), + (2, 1024, 8, 128, 1, torch.float16), + (2, 1024, 8, 128, 10, torch.float16), + (4, 2048, 8, 64, 1, torch.float16), + ] + ], +) +@pytest.mark.skipif( + device_platform == 'intel', + reason='Intel Triton Failure', +) +def test_chunk( + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + gate_logit_normalizer: float, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + # [B, T, H, D] + q = torch.rand((B, T, H, D), dtype=dtype, device=device).requires_grad_() + k = torch.rand((B, T, H, D), dtype=dtype, device=device).requires_grad_() + v = torch.rand((B, T, H, D), dtype=dtype, device=device).requires_grad_() + g = (F.logsigmoid(torch.rand((B, T, H, D), dtype=dtype, device=device)) / gate_logit_normalizer).requires_grad_() + h0 = torch.rand((B, H, D, D), dtype=dtype, device=device).requires_grad_() + do = torch.randn_like(v) + dht = torch.randn((B, H, D, D), dtype=dtype, device=device) + + tri, tri_ht = chunk_gla( + q=q, + k=k, + v=v, + g=g, + initial_state=h0, + output_final_state=True, + ) + ((tri * do).sum() + (tri_ht * dht).sum()).backward() + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_dg, g.grad = g.grad.clone(), None + tri_dh0, h0.grad = h0.grad.clone(), None + + ref, ref_ht = fused_recurrent_gla( + q=q, + k=k, + v=v, + gk=g, + initial_state=h0, + output_final_state=True, + ) + ((ref * do).sum() + (ref_ht * dht).sum()).backward() + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_dg, g.grad = g.grad.clone(), None + ref_dh0, h0.grad = h0.grad.clone(), None + + assert_close('o', ref, tri, 0.004) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) + assert_close('dg', ref_dg, tri_dg, 0.005) + assert_close('dh0', ref_dh0, tri_dh0, 0.005) + + +@pytest.mark.parametrize( + ('H', 'D', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-D{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (4, 64, [0, 15], torch.float16), + (4, 64, [0, 256, 500, 1000], torch.float16), + (4, 100, [0, 15, 100, 300, 1200, 2000], torch.float16), + ] + ], +) +@pytest.mark.skipif( + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set', +) +def test_chunk_varlen( + H: int, + D: int, + cu_seqlens: list[int], + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + + N = len(cu_seqlens) - 1 + T = cu_seqlens[-1] + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + + q = torch.rand((1, T, H, D), dtype=dtype, device=device).requires_grad_() + k = torch.rand((1, T, H, D), dtype=dtype, device=device).requires_grad_() + v = torch.rand((1, T, H, D), dtype=dtype, device=device).requires_grad_() + g = F.logsigmoid(torch.rand((1, T, H, D), dtype=dtype, device=device)).requires_grad_() + h0 = torch.rand((N, H, D, D), dtype=dtype, device=device).requires_grad_() + do = torch.randn_like(v) + dht = torch.rand((N, H, D, D), dtype=dtype, device=device) + + ref, ref_ht = fused_recurrent_gla( + q=q, + k=k, + v=v, + gk=g, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + + ((ref * do).sum() + (ref_ht * dht).sum()).backward() + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_dg, g.grad = g.grad.clone(), None + ref_dh0, h0.grad = h0.grad.clone(), None + + tri, tri_ht = chunk_gla( + q=q, + k=k, + v=v, + g=g, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + ((tri * do).sum() + (tri_ht * dht).sum()).backward() + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_dg, g.grad = g.grad.clone(), None + tri_dh0, h0.grad = h0.grad.clone(), None + + assert_close('o', ref, tri, 0.004) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) + assert_close('dg', ref_dg, tri_dg, 0.005) + assert_close('dh0', ref_dh0, tri_dh0, 0.005) diff --git a/code/flash-linear-attention/tests/ops/test_gsa.py b/code/flash-linear-attention/tests/ops/test_gsa.py new file mode 100644 index 0000000000000000000000000000000000000000..f3316a872ab605fee25f0aed2b35e134c99ab578 --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_gsa.py @@ -0,0 +1,432 @@ + +import os + +import pytest +import torch +import torch.nn.functional as F + +from fla.ops.gsa import chunk_gsa, fused_recurrent_gsa +from fla.ops.gsa.naive import naive_recurrent_gsa +from fla.utils import assert_close, check_shared_mem, device, device_platform + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'M', 'gate_logit_normalizer', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-M{}-gate_logit_normalizer{}-{}".format(*test)) + for test in [ + (1, 63, 1, 64, 32, 1, torch.float), + (2, 1024, 4, 60, 64, 1, torch.float), + (2, 1024, 8, 128, 64, 0.1, torch.float), + (2, 1024, 8, 128, 32, 1, torch.float), + (2, 1024, 8, 128, 64, 1, torch.float), + (2, 1024, 8, 128, 64, 10, torch.float), + (4, 2048, 8, 64, 64, 1, torch.float), + (2, 1024, 8, 128, 64, 0.1, torch.float16), + (2, 1024, 8, 128, 64, 10, torch.float16), + ] + ], +) +@pytest.mark.skipif( + device_platform == 'intel', + reason='Intel Triton Failure', +) +def test_fused_recurrent( + B: int, + T: int, + H: int, + D: int, + M: int, + gate_logit_normalizer: float, + dtype: torch.dtype, +): + torch.manual_seed(42) + + q = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() + k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() + v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() + s = torch.randn((B, T, H, M), dtype=dtype, device=device).requires_grad_() + g = (F.logsigmoid(torch.randn((B, T, H, M), dtype=dtype, device=device)) / gate_logit_normalizer).requires_grad_() + hk0 = torch.randn(B, H, D, M, device=device).requires_grad_() + hv0 = torch.randn(B, H, M, D, device=device).requires_grad_() + do = torch.randn_like(v) + dhkt = torch.randn_like(hk0) + dhvt = torch.randn_like(hv0) + + ref, (ref_hkt, ref_hvt) = naive_recurrent_gsa(q, k, v, s, g, initial_state=(hk0, hv0), output_final_state=True) + ((ref * do).sum() + (ref_hkt * dhkt).sum() + (ref_hvt * dhvt).sum()).backward() + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_ds, s.grad = s.grad.clone(), None + ref_dg, g.grad = g.grad.clone(), None + ref_dhk0, hk0.grad = hk0.grad.clone(), None + ref_dhv0, hv0.grad = hv0.grad.clone(), None + + tri, (tri_hkt, tri_hvt) = fused_recurrent_gsa( + q=q, + k=k, + v=v, + s=s, + g=g, + initial_state=(hk0, hv0), + output_final_state=True, + ) + ((tri * do).sum() + (tri_hkt * dhkt).sum() + (tri_hvt * dhvt).sum()).backward() + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_ds, s.grad = s.grad.clone(), None + tri_dg, s.grad = g.grad.clone(), None + tri_dhk0, hk0.grad = hk0.grad.clone(), None + tri_dhv0, hv0.grad = hv0.grad.clone(), None + + assert_close('o', ref, tri, 0.005) + assert_close('hkt', ref_hkt, tri_hkt, 0.005) + assert_close('hvt', ref_hvt, tri_hvt, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) + assert_close('ds', ref_ds, tri_ds, 0.005) + assert_close('dg', ref_dg, tri_dg, 0.005) + assert_close('dhk0', ref_dhk0, tri_dhk0, 0.005) + assert_close('dhv0', ref_dhv0, tri_dhv0, 0.005) + + +@pytest.mark.parametrize( + ('H', 'D', 'M', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-D{}-M{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (4, 64, 64, [0, 15], torch.float), + (4, 64, 64, [0, 256, 500, 1000], torch.float), + (4, 100, 64, [0, 15, 100, 300, 1200, 2000], torch.float), + (4, 64, 64, [0, 1, 100, 300, 1200, 2048], torch.float16), + (4, 128, 64, [0, 200, 512, 1200, 2048], torch.float16), + ] + ], +) +@pytest.mark.skipif( + device_platform == 'intel', + reason='Intel Triton Failure', +) +def test_fused_recurrent_varlen( + H: int, + D: int, + M: int, + cu_seqlens: list[int], + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + N = len(cu_seqlens) - 1 + T = cu_seqlens[-1] + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + + q = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + k = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + s = torch.randn((1, T, H, M), dtype=dtype, device=device).requires_grad_() + g = F.logsigmoid(torch.randn((1, T, H, M), dtype=dtype, device=device)).requires_grad_() + hk0 = torch.randn(N, H, D, M, device=device).requires_grad_() + hv0 = torch.randn(N, H, M, D, device=device).requires_grad_() + dhkt = torch.randn(N, H, D, M, device=device).requires_grad_() + dhvt = torch.randn(N, H, M, D, device=device).requires_grad_() + + do = torch.randn_like(v) + refs, ref_hkts, ref_hfts = [], [], [] + for i in range(N): + ref, (ref_hkt, ref_hvt) = naive_recurrent_gsa( + q[:, cu_seqlens[i]:cu_seqlens[i+1]], + k[:, cu_seqlens[i]:cu_seqlens[i+1]], + v[:, cu_seqlens[i]:cu_seqlens[i+1]], + s[:, cu_seqlens[i]:cu_seqlens[i+1]], + g[:, cu_seqlens[i]:cu_seqlens[i+1]], + initial_state=(hk0[i:i+1], hv0[i:i+1]), + output_final_state=True, + ) + refs.append(ref) + ref_hkts.append(ref_hkt) + ref_hfts.append(ref_hvt) + ref = torch.cat(refs, 1) + ref_hkt = torch.cat(ref_hkts, 0) + ref_hvt = torch.cat(ref_hfts, 0) + ((ref * do).sum() + (ref_hkt * dhkt).sum() + (ref_hvt * dhvt).sum()).backward() + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_ds, s.grad = s.grad.clone(), None + ref_dg, g.grad = g.grad.clone(), None + ref_dhk0, hk0.grad = hk0.grad.clone(), None + ref_dhv0, hv0.grad = hv0.grad.clone(), None + + tri, (tri_hkt, tri_hvt) = fused_recurrent_gsa( + q=q, + k=k, + v=v, + s=s, + g=g, + initial_state=(hk0, hv0), + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + ((tri * do).sum() + (tri_hkt * dhkt).sum() + (tri_hvt * dhvt).sum()).backward() + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_ds, s.grad = s.grad.clone(), None + tri_dg, s.grad = g.grad.clone(), None + tri_dhk0, hk0.grad = hk0.grad.clone(), None + tri_dhv0, hv0.grad = hv0.grad.clone(), None + + assert_close('o', ref, tri, 0.005) + assert_close('hkt', ref_hkt, tri_hkt, 0.005) + assert_close('hvt', ref_hvt, tri_hvt, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) + assert_close('ds', ref_ds, tri_ds, 0.005) + assert_close('dg', ref_dg, tri_dg, 0.005) + assert_close('dhk0', ref_dhk0, tri_dhk0, 0.005) + assert_close('dhv0', ref_dhv0, tri_dhv0, 0.005) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'M', 'gate_logit_normalizer', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-M{}-gate_logit_normalizer{}-{}".format(*test)) + for test in [ + (1, 63, 1, 64, 32, 1, torch.float16), + (2, 1024, 4, 60, 64, 1, torch.float16), + (2, 1024, 4, 256, 64, 1, torch.float16), + (2, 1024, 4, 128, 64, 0.1, torch.float), + (2, 1024, 4, 128, 128, 1, torch.float16), + (2, 1024, 4, 128, 64, 10, torch.float16), + ] + ], +) +@pytest.mark.skipif( + device_platform == 'intel', + reason='Intel Triton Failure', +) +def test_chunk( + B: int, + T: int, + H: int, + D: int, + M: int, + gate_logit_normalizer: float, + dtype: torch.dtype, +): + if (D > 64 or M > 64) and check_shared_mem('hopper') is False: + pytest.skip(reason='Current CI do not support this config') + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + + q = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() + k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() + v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() + s = torch.randn((B, T, H, M), dtype=dtype, device=device).requires_grad_() + g = (F.logsigmoid(torch.randn((B, T, H, M), dtype=dtype, device=device)) / gate_logit_normalizer).requires_grad_() + hk0 = torch.randn(B, H, D, M, device=device).requires_grad_() + hv0 = torch.randn(B, H, M, D, device=device).requires_grad_() + dhkt = torch.randn(B, H, D, M, device=device).requires_grad_() + dhvt = torch.randn(B, H, M, D, device=device).requires_grad_() + + do = torch.randn_like(v) + ref, (ref_hkt, ref_hvt) = fused_recurrent_gsa( + q=q, + k=k, + v=v, + s=s, + g=g, + scale=D**-0.5, + initial_state=(hk0, hv0), + output_final_state=True) + ((ref * do).sum() + (ref_hkt * dhkt).sum() + (ref_hvt * dhvt).sum()).backward() + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_ds, s.grad = s.grad.clone(), None + ref_dg, g.grad = g.grad.clone(), None + ref_dhk0, hk0.grad = hk0.grad.clone(), None + ref_dhv0, hv0.grad = hv0.grad.clone(), None + + tri, (tri_hkt, tri_hvt) = chunk_gsa( + q=q, + k=k, + v=v, + s=s, + g=g, + scale=D**-0.5, + initial_state=(hk0, hv0), + output_final_state=True, + ) + ((tri * do).sum() + (tri_hkt * dhkt).sum() + (tri_hvt * dhvt).sum()).backward() + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_ds, s.grad = s.grad.clone(), None + tri_dg, s.grad = g.grad.clone(), None + tri_dhk0, hk0.grad = hk0.grad.clone(), None + tri_dhv0, hv0.grad = hv0.grad.clone(), None + + assert_close('o', ref, tri, 0.005) + assert_close('hkt', ref_hkt, tri_hkt, 0.005) + assert_close('hvt', ref_hvt, tri_hvt, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) + assert_close('ds', ref_ds, tri_ds, 0.008) + assert_close('dg', ref_dg, tri_dg, 0.008) + assert_close('dhk0', ref_dhk0, tri_dhk0, 0.005) + assert_close('dhv0', ref_dhv0, tri_dhv0, 0.005) + + +@pytest.mark.parametrize( + ('H', 'D', 'M', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-D{}-M{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (4, 64, 64, [0, 15], torch.float16), + (4, 64, 64, [0, 256, 500, 1000], torch.float16), + (4, 100, 64, [0, 15, 100, 300, 1200, 2000], torch.float16), + ] + ], +) +@pytest.mark.skipif( + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set', +) +@pytest.mark.skipif( + device_platform == 'intel', + reason='Intel Triton Failure', +) +def test_chunk_varlen( + H: int, + D: int, + M: int, + cu_seqlens: list[int], + dtype: torch.dtype, +): + if (D > 64 or M > 64) and check_shared_mem('hopper') is False: + pytest.skip(reason='Current CI do not support this config') + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + N = len(cu_seqlens) - 1 + T = cu_seqlens[-1] + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + + q = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + k = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + s = torch.randn((1, T, H, M), dtype=dtype, device=device).requires_grad_() + g = F.logsigmoid(torch.randn((1, T, H, M), dtype=dtype, device=device)).requires_grad_() + hk0 = torch.randn(N, H, D, M, device=device).requires_grad_() + hv0 = torch.randn(N, H, M, D, device=device).requires_grad_() + dhkt = torch.randn(N, H, D, M, device=device).requires_grad_() + dhvt = torch.randn(N, H, M, D, device=device).requires_grad_() + + do = torch.randn_like(v) + + ref, (ref_hkt, ref_hvt) = fused_recurrent_gsa( + q=q, + k=k, + v=v, + s=s, + g=g, + scale=D**-0.5, + initial_state=(hk0, hv0), + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + ((ref * do).sum() + (ref_hkt * dhkt).sum() + (ref_hvt * dhvt).sum()).backward() + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_ds, s.grad = s.grad.clone(), None + ref_dg, g.grad = g.grad.clone(), None + ref_dhk0, hk0.grad = hk0.grad.clone(), None + ref_dhv0, hv0.grad = hv0.grad.clone(), None + + tri, (tri_hkt, tri_hvt) = chunk_gsa( + q=q, + k=k, + v=v, + s=s, + g=g, + scale=D**-0.5, + initial_state=(hk0, hv0), + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + ((tri * do).sum() + (tri_hkt * dhkt).sum() + (tri_hvt * dhvt).sum()).backward() + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_ds, s.grad = s.grad.clone(), None + tri_dg, g.grad = g.grad.clone(), None + tri_dhk0, hk0.grad = hk0.grad.clone(), None + tri_dhv0, hv0.grad = hv0.grad.clone(), None + + assert_close('o', ref, tri, 0.004) + assert_close('hkt', ref_hkt, tri_hkt, 0.005) + assert_close('hvt', ref_hvt, tri_hvt, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) + assert_close('ds', ref_ds, tri_ds, 0.005) + assert_close('dg', ref_dg, tri_dg, 0.005) + assert_close('dhk0', ref_dhk0, tri_dhk0, 0.005) + assert_close('dhv0', ref_dhv0, tri_dhv0, 0.005) + + +@pytest.mark.parametrize( + ('B', 'T', 'HQ', 'H', 'D', 'M', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-HQ{}-H{}-D{}-M{}-{}".format(*test)) + for test in [ + (2, 63, 2, 1, 64, 32, torch.float), + (2, 200, 8, 2, 64, 64, torch.float), + (2, 256, 16, 4, 128, 64, torch.float), + ] + ], +) +@pytest.mark.skipif( + device_platform == 'intel', + reason='Intel Triton Failure', +) +def test_inference( + B: int, + T: int, + HQ: int, + H: int, + D: int, + M: int, + dtype: torch.dtype, +): + torch.manual_seed(42) + + q = torch.randn((B, T, HQ, D), dtype=dtype, device=device) + k = torch.randn((B, T, H, D), dtype=dtype, device=device) + v = torch.randn((B, T, H, D), dtype=dtype, device=device) + s = torch.randn((B, T, H, M), dtype=dtype, device=device) + g = F.logsigmoid(torch.randn((B, T, H, M), dtype=dtype, device=device)) + h0 = (torch.randn(B, H, D, M, dtype=dtype, device=device), + torch.randn(B, H, M, D, dtype=dtype, device=device)) + + ref, _ = naive_recurrent_gsa(q, k, v, s, g, initial_state=h0) + tri = torch.empty_like(ref) + for i in range(T): + o, ht = fused_recurrent_gsa( + q[:, i:i+1], + k[:, i:i+1], + v[:, i:i+1], + s[:, i:i+1], + g[:, i:i+1], + initial_state=h0, + output_final_state=True, + ) + tri[:, i] = o.squeeze(1) + assert_close(f'o{i}', ref[:, i], tri[:, i], 0.005) + h0 = ht diff --git a/code/flash-linear-attention/tests/ops/test_hgrn.py b/code/flash-linear-attention/tests/ops/test_hgrn.py new file mode 100644 index 0000000000000000000000000000000000000000..cafb4d400a0ce22fa5dcbdb7446727919acfac35 --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_hgrn.py @@ -0,0 +1,164 @@ + +import os + +import pytest +import torch +import torch.nn.functional as F + +from fla.ops.hgrn import chunk_hgrn, fused_recurrent_hgrn +from fla.ops.hgrn.naive import naive_recurrent_hgrn +from fla.utils import assert_close, device + + +@pytest.mark.parametrize( + ('B', 'T', 'D', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-D{}-{}".format(*test)) + for test in [ + (1, 63, 500, torch.float), + (2, 1024, 500, torch.float), + (2, 1024, 512, torch.float), + (2, 1024, 1000, torch.float), + (4, 2048, 2048, torch.float), + ] + ], +) +def test_fused_recurrent( + B: int, + T: int, + D: int, + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + + x = torch.randn((B, T, D), dtype=dtype, device=device) + g = torch.randn((B, T, D), dtype=dtype, device=device) + h0 = torch.randn_like(x[:, 0]) + x, g = (1 - g.sigmoid()) * x, F.logsigmoid(g) + x, g, h0 = (i.detach().clone().to(dtype).requires_grad_() for i in (x, g, h0)) + + do = torch.randn_like(x) + dht = torch.randn_like(h0) + ref, ref_ht = naive_recurrent_hgrn(x, g, h0, output_final_state=True) + ((ref * do).sum() + (ref_ht * dht).sum()).backward() + ref_dx, x.grad = x.grad.clone(), None + ref_dg, g.grad = g.grad.clone(), None + ref_dh0, h0.grad = h0.grad.clone(), None + + tri, tri_ht = fused_recurrent_hgrn(x, g, h0, output_final_state=True) + ((tri * do).sum() + (tri_ht * dht).sum()).backward() + tri_dx, x.grad = x.grad.clone(), None + tri_dg, g.grad = g.grad.clone(), None + tri_dh0, h0.grad = h0.grad.clone(), None + + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dx', ref_dx, tri_dx, 0.005) + assert_close('dg', ref_dg, tri_dg, 0.005) + assert_close('dh0', ref_dh0, tri_dh0, 0.005) + + +@pytest.mark.parametrize( + ('D', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="D{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (500, [0, 15], torch.float), + (512, [0, 256, 500, 1000], torch.float), + (1000, [0, 15, 100, 300, 1200, 2000], torch.float), + (2048, [0, 200, 512, 1200, 2048], torch.float16), + ] + ], +) +def test_fused_recurrent_varlen( + D: int, + cu_seqlens: list[int], + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + + N = len(cu_seqlens) - 1 + T = cu_seqlens[-1] + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + + x = torch.randn((1, T, D), dtype=dtype, device=device) + g = torch.randn((1, T, D), dtype=dtype, device=device) + h0 = torch.randn(N, D, dtype=dtype, device=device) + x, g = (1 - g.sigmoid()) * x, F.logsigmoid(g) + x, g, h0 = (i.detach().clone().to(dtype).requires_grad_() for i in (x, g, h0)) + + do = torch.randn_like(x) + dht = torch.randn_like(h0) + refs, ref_hts = [], [] + for i in range(N): + ref, ref_ht = naive_recurrent_hgrn( + x[:, cu_seqlens[i]:cu_seqlens[i+1]], + g[:, cu_seqlens[i]:cu_seqlens[i+1]], + h0[i:i+1], + output_final_state=True, + ) + refs.append(ref) + ref_hts.append(ref_ht) + ref = torch.cat(refs, 1) + ref_ht = torch.cat(ref_hts, 0) + ((ref * do).sum() + (ref_ht * dht).sum()).backward() + ref_dx, x.grad = x.grad.clone(), None + ref_dg, g.grad = g.grad.clone(), None + ref_dh0, h0.grad = h0.grad.clone(), None + + tri, tri_ht = fused_recurrent_hgrn(x, g, h0, output_final_state=True, cu_seqlens=cu_seqlens) + ((tri * do).sum() + (tri_ht * dht).sum()).backward() + tri_dx, x.grad = x.grad.clone(), None + tri_dg, g.grad = g.grad.clone(), None + tri_dh0, h0.grad = h0.grad.clone(), None + + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dx', ref_dx, tri_dx, 0.005) + assert_close('dg', ref_dg, tri_dg, 0.005) + assert_close('dh0', ref_dh0, tri_dh0, 0.005) + + +@pytest.mark.parametrize( + ('B', 'T', 'D', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-D{}-{}".format(*test)) + for test in [ + (1, 63, 500, torch.float16), + (2, 500, 1000, torch.float16), + (2, 1000, 1024, torch.float16), + (4, 2048, 2048, torch.float16), + ] + ], +) +def test_chunk( + B: int, + T: int, + D: int, + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + + x = torch.randn((B, T, D), dtype=dtype, device=device) + g = torch.randn((B, T, D), dtype=dtype, device=device) + x, g = (1 - g.sigmoid()) * x, F.logsigmoid(g) + x, g = (i.detach().clone().to(dtype).requires_grad_() for i in (x, g)) + + do = torch.randn_like(x) + h0 = torch.randn_like(x[:, 0]) + ref, _ = fused_recurrent_hgrn(x, g, h0, output_final_state=True) + ref.backward(do) + ref_dx, x.grad = x.grad.clone(), None + ref_dg, g.grad = g.grad.clone(), None + + tri, _ = chunk_hgrn(x, g, h0, output_final_state=True) + tri.backward(do) + tri_dx, x.grad = x.grad.clone(), None + tri_dg, g.grad = g.grad.clone(), None + + assert_close('o', ref, tri, 0.005) + assert_close('dx', ref_dx, tri_dx, 0.005) + assert_close('dg', ref_dg, tri_dg, 0.005) diff --git a/code/flash-linear-attention/tests/ops/test_iplr_delta.py b/code/flash-linear-attention/tests/ops/test_iplr_delta.py new file mode 100644 index 0000000000000000000000000000000000000000..7e6e6240b281234cb8959e8b8fe57829ed8e208f --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_iplr_delta.py @@ -0,0 +1,235 @@ + + +import pytest +import torch +import torch.nn.functional as F +from einops import rearrange + +from fla.ops.generalized_delta_rule.iplr.chunk import chunk_iplr_delta_rule +from fla.ops.generalized_delta_rule.iplr.fused_recurrent import fused_recurrent_iplr_delta_rule +from fla.utils import assert_close, device + + +def chunk_iplr_delta_rule_ref( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + initial_state: torch.Tensor = None, + output_final_state: bool = True, + scale: float = None, + chunk_size: int = 64, +): + BT = chunk_size + if scale is None: + scale = 1 / (q.shape[-1] ** 0.5) + + q, k, v, a, b = map(lambda x: x.transpose(1, 2), (q, k, v, a, b)) + T = q.shape[-2] + pad_len = (BT - (T % BT)) % BT + if pad_len > 0: + q = F.pad(q, (0, 0, 0, pad_len)) + k = F.pad(k, (0, 0, 0, pad_len)) + v = F.pad(v, (0, 0, 0, pad_len)) + a = F.pad(a, (0, 0, 0, pad_len)) + b = F.pad(b, (0, 0, 0, pad_len)) + q, k, v, a, b = map(lambda x: x.to(torch.float32), [q, k, v, a, b]) + + B, H, L, DK = q.shape + DV = v.shape[-1] + q = q * scale + + S = k.new_zeros(B, H, DK, DV) + if initial_state is not None: + S += initial_state + + # note that diagonal is masked. + mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0) + q, k, v, a, b = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size), [q, k, v, a, b]) + + v2 = (a @ k.transpose(-1, -2)).masked_fill_(mask, 0) @ v + attn = (a @ b.transpose(-1, -2)).masked_fill(mask, 0) + for i in range(1, chunk_size): + attn[..., i, :i] = attn[..., i, :i] + (attn[..., i, :, None].clone() * attn[..., :, :i].clone()).sum(-2) + attn = attn + torch.eye(chunk_size, dtype=torch.float, device=q.device) + u = attn @ v2 + w = attn @ a + o = torch.zeros_like(v) + mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1) + for i in range(0, L // chunk_size): + current_chunk_size = min(chunk_size, L - i * chunk_size) # to handle the last chunk with possibly padding + q_i = q[:, :, i, :current_chunk_size] + k_i = k[:, :, i, :current_chunk_size] + v_i = v[:, :, i, :current_chunk_size] + u_i = u[:, :, i, :current_chunk_size] + w_i = w[:, :, i, :current_chunk_size] + b_i = b[:, :, i, :current_chunk_size] + o_1 = (q_i @ k_i.transpose(-1, -2)).masked_fill_(mask, 0) @ v_i + v2_i = u_i + w_i @ S + o_2 = (q_i @ b_i.transpose(-1, -2)).masked_fill_(mask, 0) @ v2_i + o_3 = q_i @ S + o[:, :, i, :current_chunk_size] = o_1 + o_2 + o_3 + S = S + k_i.transpose(-1, -2) @ v_i + b_i.transpose(-1, -2) @ v2_i + S = None if output_final_state is False else S + o = rearrange(o, 'b h n c d -> b h (n c) d') + o = o[:, :, :T] + o = o.transpose(1, 2) + return o, S + + +def recurrence_iplr_delta_rule_ref( + q, + k, + v, + a, + b, + initial_state: torch.Tensor | None = None, + output_final_state: bool = True, + scale: float | None = None, +): + orig_dtype = q.dtype + if scale is None: + scale = 1 / (q.shape[-1] ** 0.5) + q, k, v, a, b = map(lambda x: x.transpose(1, 2).to(torch.float32), [q, k, v, a, b]) + q = q * scale + B, H, L, DK = q.shape + DV = v.shape[-1] + o = torch.zeros_like(v) + S = torch.zeros(B, H, DK, DV).to(v) + if initial_state is not None: + S += initial_state + + for i in range(q.shape[-2]): + _k = k[:, :, i] + _q = q[:, :, i] + _v = v[:, :, i] + _a = a[:, :, i] + _b = b[:, :, i] + _kv = _k[..., None] * _v[..., None, :] + (S.clone() * _a[..., None]).sum(-2, keepdim=True) * _b[..., None] + S = S + _kv + o[:, :, i] = torch.einsum('bhd,bhdm->bhm', _q, S) + S = None if output_final_state is False else S + o = o.transpose(1, 2) + return o.to(orig_dtype), S + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'scale', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-{}".format(*test)) + for test in [ + (1, 63, 1, 64, 1, torch.float), + (2, 1024, 4, 60, 1, torch.float), + (2, 1024, 8, 100, 1, torch.float), + (2, 1024, 8, 128, 0.1, torch.float), + (4, 2048, 8, 64, 0.1, torch.float), + ] + ], +) +def test_fused_recurrent( + B: int, + T: int, + H: int, + D: int, + scale: float, + dtype: torch.dtype, +): + q = torch.randn(B, T, H, D, dtype=dtype) + k = torch.randn(B, T, H, D, dtype=dtype) + v = torch.randn(B, T, H, D, dtype=dtype) + a = torch.rand(B, T, H, D, dtype=dtype) + + a = F.normalize(a, p=2, dim=-1) + b = -a + h0 = torch.zeros(B, H, D, D, dtype=torch.float32) + q, k, v, a, b, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, a, b, h0)) + ref, ref_ht = recurrence_iplr_delta_rule_ref( + q=q.clone(), + k=k.clone(), + v=v.clone(), + a=a.clone(), + b=b.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + dht = torch.rand_like(h0) + do = torch.rand_like(ref) + ((dht * ref_ht).sum() + (do * ref).sum()).backward() + dq, dk, dv, da, db, dh0 = map(lambda x: x.grad, (q, k, v, a, b, h0)) + q.grad, k.grad, v.grad, a.grad, b.grad, h0.grad = None, None, None, None, None, None + tri, tri_ht = fused_recurrent_iplr_delta_rule( + q=q.clone(), + k=k.clone(), + v=v.clone(), + a=a.clone(), + b=b.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + ((dht * tri_ht).sum() + (do * tri).sum()).backward() + assert_close('o', ref, tri, 0.003) + assert_close('ht', ref_ht, tri_ht, 0.003) + assert_close('dq', dq, q.grad, 0.003) + assert_close('dk', dk, k.grad, 0.003) + assert_close('dv', dv, v.grad, 0.003) + assert_close('da', da, a.grad, 0.003) + assert_close('db', db, b.grad, 0.003) + assert_close('dh0', dh0, h0.grad, 0.003) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'scale', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-{}".format(*test)) + for test in [ + (1, 63, 1, 64, 1, torch.float16), + (2, 500, 3, 60, 1, torch.float16), + (2, 1000, 3, 64, 0.1, torch.float16), + (2, 1024, 4, 100, 1, torch.float16), + (3, 1024, 4, 128, 0.1, torch.float16), + (4, 2048, 8, 64, 0.1, torch.float16), + ] + ], +) +def test_chunk( + B: int, + T: int, + H: int, + D: int, + scale: float, + dtype: torch.dtype, +): + q = torch.randn(B, T, H, D, dtype=dtype) + k = torch.randn(B, T, H, D, dtype=dtype) + v = torch.randn(B, T, H, D, dtype=dtype) + a = torch.rand(B, T, H, D, dtype=dtype) + + a = F.normalize(a, p=2, dim=-1) + b = -a + h0 = torch.zeros(B, H, D, D, dtype=torch.float32) + q, k, v, a, b, h0 = map(lambda x: x.to(device).requires_grad_(), (q, k, v, a, b, h0)) + ref, ref_ht = recurrence_iplr_delta_rule_ref( + q=q.clone(), + k=k.clone(), + v=v.clone(), + a=a.clone(), + b=b.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + tri, tri_ht = chunk_iplr_delta_rule( + q=q.clone(), + k=k.clone(), + v=v.clone(), + a=a.clone(), + b=b.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + assert_close('o', ref, tri, 0.007) + assert_close('ht', ref_ht, tri_ht, 0.008) diff --git a/code/flash-linear-attention/tests/ops/test_kda.py b/code/flash-linear-attention/tests/ops/test_kda.py new file mode 100644 index 0000000000000000000000000000000000000000..3a2851207a7c6cbf6c3a9d860508f1c3ab16bc54 --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_kda.py @@ -0,0 +1,379 @@ + +import os + +import pytest +import torch +import torch.nn.functional as F + +from fla.ops.kda import chunk_kda, fused_recurrent_kda +from fla.ops.kda.gate import fused_kda_gate, kda_gate_ref +from fla.ops.kda.naive import naive_chunk_kda, naive_recurrent_kda +from fla.utils import assert_close, device, is_intel_alchemist + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'dtype'), + [ + pytest.param( + *test, + id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-{}".format(*test), + ) + for test in [ + (1, 64, 1, 64, 1, 1, torch.float), + (2, 512, 3, 60, 1, 1, torch.float), + (4, 1024, 4, 128, 0.1, 1, torch.float), + (4, 1024, 4, 128, 1, 10, torch.float), + ] + ], +) +def test_naive_chunk( + B: int, + T: int, + H: int, + D: int, + scale: float, + gate_logit_normalizer: float, + dtype: torch.dtype, +): + torch.manual_seed(42) + if is_intel_alchemist and D > 128: + pytest.skip(reason='chunk_gated_delta_rule is not supported on alchemist for D>128') + + q = torch.rand(B, T, H, D, dtype=dtype) + k = torch.rand(B, T, H, D, dtype=dtype) + v = torch.rand(B, T, H, D, dtype=dtype) + g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float)) / gate_logit_normalizer + beta = torch.randn(B, T, H, dtype=dtype).sigmoid() + h0 = torch.randn(B, H, D, D, dtype=torch.float32) + q, k, v, g, beta, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, g, beta, h0)) + + ref, ref_ht = naive_recurrent_kda( + q=F.normalize(q.clone(), p=2, dim=-1), + k=F.normalize(k.clone(), p=2, dim=-1), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + + tri, tri_ht = naive_chunk_kda( + q=F.normalize(q.clone(), p=2, dim=-1), + k=F.normalize(k.clone(), p=2, dim=-1), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'use_qk_l2norm_in_kernel', 'dtype'), + [ + pytest.param( + *test, + id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-use_qk_l2norm_in_kernel{}-{}".format(*test), + ) + for test in [ + (1, 64, 1, 64, 1, 1, False, torch.float), + (2, 512, 3, 60, 1, 1, False, torch.float), + (3, 1000, 4, 100, 0.1, 1, True, torch.float), + (4, 1024, 4, 128, 0.1, 1, False, torch.float), + ] + ], +) +def test_fused_recurrent( + B: int, + T: int, + H: int, + D: int, + scale: float, + gate_logit_normalizer: float, + use_qk_l2norm_in_kernel: bool, + dtype: torch.dtype, +): + torch.manual_seed(42) + if is_intel_alchemist and D > 128: + pytest.skip(reason='chunk_gated_delta_rule is not supported on alchemist for D>128') + + q = torch.rand(B, T, H, D, dtype=dtype) + k = torch.rand(B, T, H, D, dtype=dtype) + v = torch.rand(B, T, H, D, dtype=dtype) + g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float)) / gate_logit_normalizer + beta = torch.randn(B, T, H, dtype=dtype).sigmoid() + h0 = torch.randn(B, H, D, D, dtype=torch.float32) + q, k, v, g, beta, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, g, beta, h0)) + + ref, ref_ht = naive_recurrent_kda( + q=F.normalize(q.clone(), p=2, dim=-1), + k=F.normalize(k.clone(), p=2, dim=-1), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + + tri, tri_ht = fused_recurrent_kda( + q=F.normalize(q.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q.clone(), + k=F.normalize(k.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k.clone(), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'mask_p', 'use_qk_l2norm_in_kernel', 'dtype', 'tma'), + [ + pytest.param( + *test, + id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-mask_p{}-use_qk_l2norm_in_kernel{}-{}-tma{}".format(*test), + ) + for test in [ + (1, 63, 1, 64, 1, 1, 0, False, torch.float16, True), + (2, 500, 3, 60, 1, 1, 0, False, torch.float16, True), + (2, 1000, 3, 64, 0.1, 1, 0.5, False, torch.float16, False), + (3, 1024, 4, 100, 1, 0.1, 0, False, torch.float16, False), + (4, 1024, 4, 128, 0.1, 1, 0, False, torch.float16, True), + (4, 1024, 4, 128, 0.1, 1, 0, True, torch.float16, True), + (2, 1500, 4, 128, 0.1, 10, 0, False, torch.float16, False), + (4, 2048, 8, 64, 0.1, 1, 0, False, torch.float16, True), + ] + ], +) +def test_chunk( + B: int, + T: int, + H: int, + D: int, + scale: float, + gate_logit_normalizer: float, + mask_p: float, + use_qk_l2norm_in_kernel: bool, + dtype: torch.dtype, + tma: bool, +): + torch.manual_seed(42) + if not tma: + os.environ['FLA_USE_TMA'] = '0' + else: + os.environ['FLA_USE_TMA'] = '1' + q = torch.rand(B, T, H, D, dtype=dtype) + k = torch.rand(B, T, H, D, dtype=dtype) + v = torch.rand(B, T, H, D, dtype=dtype) + g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float)) / gate_logit_normalizer + g = g * (torch.rand_like(g) > mask_p) + beta = torch.randn(B, T, H, dtype=dtype).sigmoid() + h0 = torch.randn(B, H, D, D, dtype=torch.float32) + q, k, v, g, beta, h0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, g, beta, h0)) + do = torch.randn_like(v) + dht = torch.randn_like(h0) + + ref, ref_ht = naive_recurrent_kda( + q=F.normalize(q.clone(), p=2, dim=-1), + k=F.normalize(k.clone(), p=2, dim=-1), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) + ref_dq, ref_dk, ref_dv, ref_dg, ref_db, ref_dh0 = q.grad, k.grad, v.grad, g.grad, beta.grad, h0.grad + q.grad = k.grad = v.grad = g.grad = beta.grad = h0.grad = None + + tri, tri_ht = chunk_kda( + q=F.normalize(q.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else q.clone(), + k=F.normalize(k.clone(), p=2, dim=-1) if not use_qk_l2norm_in_kernel else k.clone(), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) + ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True) + tri_dq, tri_dk, tri_dv, tri_dg, tri_db, tri_dh0 = q.grad, k.grad, v.grad, g.grad, beta.grad, h0.grad + q.grad = k.grad = v.grad = g.grad = beta.grad = h0.grad = None + + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.008) + assert_close('dk', ref_dk, tri_dk, 0.008) + assert_close('dv', ref_dv, tri_dv, 0.008) + assert_close('dg', ref_dg, tri_dg, 0.02) + assert_close('db', ref_db, tri_db, 0.02) + assert_close('dh0', ref_dh0, tri_dh0, 0.008) + + +@pytest.mark.parametrize( + ('H', 'D', 'mask_p', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-D{}-mask_p{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (4, 60, 0, [0, 15], torch.float16), + (4, 64, 0, [0, 256, 500, 1000], torch.float16), + (4, 128, 0.5, [0, 256, 500, 1000], torch.float16), + (4, 100, 0, [0, 15, 100, 300, 1200, 2000], torch.float16), + (4, 256, 0, [0, 15, 100, 300, 1200, 4096], torch.float16), + ] + ], +) +@pytest.mark.skipif( + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set', +) +def test_chunk_varlen( + H: int, + D: int, + mask_p: float, + cu_seqlens: list[int], + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + # randomly split the sequence into N segments + cu_seqlens = torch.LongTensor(cu_seqlens).to(device) + T = cu_seqlens[-1] + N = len(cu_seqlens) - 1 + + # seq-first required for inputs with variable lengths + q = torch.randn((1, T, H, D), dtype=dtype) + k = F.normalize(torch.randn(1, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype) + v = torch.randn((1, T, H, D), dtype=dtype) + g = F.logsigmoid(torch.randn(1, T, H, D, dtype=torch.float)) + g = g * (torch.rand_like(g) > mask_p) + beta = torch.rand(1, T, H, dtype=dtype).sigmoid() + h0 = torch.randn((N, H, D, D), dtype=dtype) + + q, k, v, g, beta, h0 = map(lambda x: x.to(device).requires_grad_(), (q, k, v, g, beta, h0)) + do = torch.randn_like(v) + dht = torch.rand_like(h0) + + tri, tri_ht = chunk_kda( + q=q.clone(), + k=k.clone(), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + initial_state=h0.clone(), + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + ((tri * do).sum() + (tri_ht * dht).sum()).backward(retain_graph=True) + tri_dq, tri_dk, tri_dv, tri_dg, tri_db, tri_dh0 = q.grad, k.grad, v.grad, g.grad, beta.grad, h0.grad + q.grad = k.grad = v.grad = g.grad = beta.grad = h0.grad = None + + ref = [] + ref_ht = [] + for i in range(N): + ref_i, ref_ht_i = naive_recurrent_kda( + q=q[:, cu_seqlens[i]:cu_seqlens[i+1]], + k=k[:, cu_seqlens[i]:cu_seqlens[i+1]], + v=v[:, cu_seqlens[i]:cu_seqlens[i+1]], + beta=beta[:, cu_seqlens[i]:cu_seqlens[i+1]], + g=g[:, cu_seqlens[i]:cu_seqlens[i+1]], + initial_state=h0[i], + output_final_state=True, + ) + ref.append(ref_i) + ref_ht.append(ref_ht_i) + ref = torch.cat(ref, 1) + ref_ht = torch.cat(ref_ht, 0) + + ((ref * do).sum() + (ref_ht * dht).sum()).backward(retain_graph=True) + ref_dq, ref_dk, ref_dv, ref_dg, ref_db, ref_dh0 = q.grad, k.grad, v.grad, g.grad, beta.grad, h0.grad + + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.007) + assert_close('dk', ref_dk, tri_dk, 0.008) + assert_close('dv', ref_dv, tri_dv, 0.007) + assert_close('dg', ref_dg, tri_dg, 0.015) + assert_close('db', ref_db, tri_db, 0.015) + assert_close('dh0', ref_dh0, tri_dh0, 0.007) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'use_bias'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-bias{}".format(*test)) + for test in [ + (1, 2, 2, 12, False), + (1, 32, 2, 16, False), + (2, 64, 4, 32, False), + (4, 128, 8, 64, False), + (4, 128, 8, 128, False), + # Add bias tests + (1, 2, 2, 12, True), + (1, 32, 2, 16, True), + (2, 64, 4, 32, True), + (4, 128, 8, 64, True), + (4, 128, 8, 128, True), + ] + ], +) +def test_kda_gate( + B: int, + T: int, + H: int, + D: int, + use_bias: bool, +): + """Test kda gate forward and backward pass - reference vs Triton implementation""" + torch.manual_seed(42) + + g = torch.randn(B, T, H * D, dtype=torch.float32) + # Ensure some values are > 20 to test the threshold logic in softplus + g = g * 30 # Scale up to get values > 20 + A = torch.log(torch.randn(1, 1, H, 1, dtype=torch.float32).uniform_(1, 16)) + g_bias = torch.randn(H * D, dtype=torch.float32) if use_bias else None + + # Move to device and set requires_grad + g, A = map(lambda x: x.to(device).requires_grad_(True), (g, A)) + if g_bias is not None: + g_bias = g_bias.to(device).requires_grad_(True) + + # Create gradient output + do = torch.randn_like(g).view(B, T, H, D) + + # Reference implementation + ref = kda_gate_ref(g.clone(), A.clone(), D, g_bias.clone() if g_bias is not None else None) + # Triton implementation + tri = fused_kda_gate(g.clone(), A.clone(), D, g_bias.clone() if g_bias is not None else None) + + # Backward pass + ((ref * do).sum()).backward(retain_graph=True) + ref_dg, ref_dA = g.grad, A.grad + ref_dgbias = g_bias.grad if g_bias is not None else None + g.grad = A.grad = None + if g_bias is not None: + g_bias.grad = None + + ((tri * do).sum()).backward(retain_graph=True) + tri_dg, tri_dA = g.grad, A.grad + tri_dgbias = g_bias.grad if g_bias is not None else None + g.grad = A.grad = None + if g_bias is not None: + g_bias.grad = None + + assert_close('o', ref, tri, 1e-4) + assert_close('dg', ref_dg, tri_dg, 1e-4) + assert_close('dA', ref_dA, tri_dA, 1e-4) + if use_bias: + assert_close('dgbias', ref_dgbias, tri_dgbias, 1e-4) diff --git a/code/flash-linear-attention/tests/ops/test_linear_attn.py b/code/flash-linear-attention/tests/ops/test_linear_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..cda197920a82a72c141e1d59df844655fed20a15 --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_linear_attn.py @@ -0,0 +1,189 @@ + + +import pytest +import torch + +from fla.ops.linear_attn import chunk_linear_attn, fused_chunk_linear_attn, fused_recurrent_linear_attn +from fla.ops.linear_attn.naive import naive_recurrent_linear_attn +from fla.utils import assert_close, device + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'scale', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-{}".format(*test)) + for test in [ + (1, 64, 1, 64, None, torch.float), + (2, 512, 4, 60, None, torch.float), + (3, 1024, 8, 128, 1., torch.float), + (3, 1024, 8, 128, 0.1, torch.float), + (3, 1024, 8, 128, None, torch.float), + (2, 2048, 8, 256, None, torch.float16), + (2, 2048, 4, 256, None, torch.float16), + ] + ], +) +def test_fused_recurrent( + B: int, + T: int, + H: int, + D: int, + scale: float | None, + dtype: torch.dtype, +): + torch.manual_seed(42) + q = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() + k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() + v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() + h0 = torch.randn((B, H, D, D), dtype=torch.float, device=device).requires_grad_() + do = torch.randn_like(v) + dht = torch.randn_like(h0) + + ref, ref_ht = naive_recurrent_linear_attn(q, k, v, scale=scale, initial_state=h0, output_final_state=True, normalize=False) + ((ref * do).sum() + (ref_ht * dht).sum()).backward() + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_dh0, h0.grad = h0.grad.clone(), None + + tri, tri_ht = fused_recurrent_linear_attn(q, k, v, scale=scale, initial_state=h0, output_final_state=True, normalize=False) + ((tri * do).sum() + (tri_ht * dht).sum()).backward() + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_dh0, h0.grad = h0.grad.clone(), None + + assert_close('o', ref, tri, 0.001) + assert_close('ht', ref_ht, tri_ht, 0.001) + assert_close('dq', ref_dq, tri_dq, 0.001) + assert_close('dk', ref_dk, tri_dk, 0.001) + assert_close('dv', ref_dv, tri_dv, 0.001) + assert_close('dh0', ref_dh0, tri_dh0, 0.001) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (1, 63, 1, 64, torch.float16), + (2, 500, 3, 60, torch.float16), + (2, 1000, 3, 128, torch.float16), + (3, 1000, 4, 64, torch.float16), + (2, 2048, 4, 256, torch.float16), + ] + ], +) +def test_chunk( + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): + torch.manual_seed(42) + q = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() + k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() + v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() + h0 = torch.randn((B, H, D, D), dtype=torch.float, device=device).requires_grad_() + do = torch.randn_like(v) + dht = torch.randn_like(h0) + + ref, ref_ht = fused_recurrent_linear_attn( + q.to(torch.float32), + k.to(torch.float32), + v.to(torch.float32), + initial_state=h0, + output_final_state=True, + normalize=False, + ) + ((ref * do).sum() + (ref_ht * dht).sum()).backward() + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_dh0, h0.grad = h0.grad.clone(), None + + tri, tri_ht = chunk_linear_attn( + q=q, + k=k, + v=v, + initial_state=h0, + output_final_state=True, + normalize=False, + ) + ((tri * do).sum() + (tri_ht * dht).sum()).backward() + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_dh0, h0.grad = h0.grad.clone(), None + + assert_close('o', ref, tri, 0.001) + assert_close('ht', ref_ht, tri_ht, 0.001) + assert_close('dq', ref_dq, tri_dq, 0.001) + assert_close('dk', ref_dk, tri_dk, 0.001) + assert_close('dv', ref_dv, tri_dv, 0.001) + assert_close('dh0', ref_dh0, tri_dh0, 0.001) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (1, 63, 1, 64, torch.float16), + (2, 500, 3, 60, torch.float16), + (2, 1000, 3, 128, torch.float16), + (3, 1000, 4, 64, torch.float16), + (2, 2048, 4, 256, torch.float16), + ] + ], +) +def test_fused_chunk( + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): + torch.manual_seed(42) + q = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() + k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() + v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() + h0 = torch.randn((B, H, D, D), dtype=torch.float, device=device).requires_grad_() + do = torch.randn_like(v) + dht = torch.randn_like(h0) + + ref, ref_ht = fused_recurrent_linear_attn( + q.to(torch.float32), + k.to(torch.float32), + v.to(torch.float32), + initial_state=h0, + output_final_state=True, + normalize=False, + ) + ((ref * do).sum() + (ref_ht * dht).sum()).backward() + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_dh0, h0.grad = h0.grad.clone(), None + + tri, tri_ht = fused_chunk_linear_attn( + q=q, + k=k, + v=v, + initial_state=h0, + output_final_state=True, + normalize=False, + ) + ((tri * do).sum() + (tri_ht * dht).sum()).backward() + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_dh0, h0.grad = h0.grad.clone(), None + + assert_close('o', ref, tri, 0.001) + assert_close('ht', ref_ht, tri_ht, 0.001) + assert_close('dq', ref_dq, tri_dq, 0.001) + assert_close('dk', ref_dk, tri_dk, 0.001) + assert_close('dv', ref_dv, tri_dv, 0.001) + assert_close('dh0', ref_dh0, tri_dh0, 0.001) diff --git a/code/flash-linear-attention/tests/ops/test_log_linear_attn.py b/code/flash-linear-attention/tests/ops/test_log_linear_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..6cd5b3dc0a2eacd483f0b4d11434f5700b237e22 --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_log_linear_attn.py @@ -0,0 +1,148 @@ +import os + +import numpy as np +import pytest +import torch + +from fla.ops.log_linear_attn import chunk_log_linear_attn +from fla.ops.log_linear_attn.naive import naive_log_linear_attn +from fla.utils import assert_close, device, device_platform + + +@pytest.mark.parametrize( + ("B", "T", "H", "D", "dtype"), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [(2, 1024, 8, 128, torch.float32), (4, 2048, 8, 64, torch.float32)] + ], +) +@pytest.mark.skipif(device_platform == "intel", reason="Intel Triton Failure") +def test_chunk( + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ["TRITON_F32_DEFAULT"] = "ieee" + + L = int(np.log2(T) + 1) + x = torch.randn(B, T, H, D, dtype=dtype, device=device) + dt = torch.nn.functional.softplus( + torch.randn(B, T, H, dtype=torch.float32, device=device) - 4, + ) + a = -torch.exp(torch.rand(H, dtype=torch.float32, device=device)) + q = torch.randn(B, T, 1, D, dtype=dtype, device=device) + k = torch.randn(B, T, 1, D, dtype=dtype, device=device) + level_scales = torch.randn(B, T, H, L, dtype=dtype, device=device) + v = (x * dt.unsqueeze(-1)).to(dtype=dtype) + g = a * dt + + out, _ = chunk_log_linear_attn(q, k, v, g, level_scales) + + ref = naive_log_linear_attn(q, k, v, g, level_scales) + + assert_close("o", ref, out, 0.004) + + +@pytest.mark.parametrize( + ("B", "T", "H", "D", "dtype"), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [(2, 512, 8, 64, torch.float32), (2, 1024, 8, 128, torch.float32)] + ], +) +@pytest.mark.skipif(device_platform == "intel", reason="Intel Triton Failure") +def test_chunk_bwd( + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ["TRITON_F32_DEFAULT"] = "ieee" + + L = int(np.log2(T) + 1) + x = torch.randn(B, T, H, D, dtype=dtype, device=device) + dt = torch.nn.functional.softplus( + torch.randn(B, T, H, dtype=torch.float32, device=device) - 4, + ) + a = -torch.exp(torch.rand(H, dtype=torch.float32, device=device)) + q = torch.randn(B, T, 1, D, dtype=dtype, device=device) + k = torch.randn(B, T, 1, D, dtype=dtype, device=device) + level_scales = torch.randn(B, T, H, L, dtype=dtype, device=device) + v = (x * dt.unsqueeze(-1)).to(dtype=dtype) + g = a * dt + do = torch.randn_like(v) + q, k, v, g, level_scales = map(lambda x: x.to(device).requires_grad_(), (q, k, v, g, level_scales)) + + out, _ = chunk_log_linear_attn(q, k, v, g, level_scales) + (out * do).sum().backward() + tri_dq, tri_dk, tri_dv, tri_dg, tri_dl = q.grad, k.grad, v.grad, g.grad, level_scales.grad + q.grad = k.grad = v.grad = g.grad = level_scales.grad = None + + ref = naive_log_linear_attn(q, k, v, g, level_scales) + (ref * do).sum().backward() + ref_dq, ref_dk, ref_dv, ref_dg, ref_dl = q.grad, k.grad, v.grad, g.grad, level_scales.grad + + assert_close("o", ref, out, 0.004) + assert_close("dq", ref_dq, tri_dq, 0.007) + assert_close("dk", ref_dk, tri_dk, 0.008) + assert_close("dv", ref_dv, tri_dv, 0.007) + assert_close("dg", ref_dg, tri_dg, 0.015) + assert_close("dl", ref_dl, tri_dl, 0.015) + + +@pytest.mark.parametrize( + ("H", "D", "cu_seqlens", "dtype"), + [ + pytest.param(*test, id="H{}-D{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (4, 64, [0, 15], torch.float32), + (4, 64, [0, 256, 500, 1000], torch.float32), + (4, 128, [0, 15, 100, 300, 1200, 2000], torch.float32), + ] + ], +) +@pytest.mark.skipif(device_platform == "intel", reason="Intel Triton Failure") +def test_chunk_varlen( + H: int, + D: int, + cu_seqlens: list[int], + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ["TRITON_F32_DEFAULT"] = "ieee" + + cu_seqlens = torch.LongTensor(cu_seqlens).to(device) + T = cu_seqlens[-1].item() + + L = int(np.ceil(np.log2(T)) + 1) + x = torch.randn(1, T, H, D, dtype=dtype, device=device) + dt = torch.nn.functional.softplus( + torch.randn(1, T, H, dtype=torch.float32, device=device) - 4, + ) + a = -torch.exp(torch.rand(H, dtype=torch.float32, device=device)) + q = torch.randn(1, T, 1, D, dtype=dtype, device=device) + k = torch.randn(1, T, 1, D, dtype=dtype, device=device) + level_scales = torch.randn(1, T, H, L, dtype=dtype, device=device) + v = (x * dt.unsqueeze(-1)).to(dtype=dtype) + g = a * dt + + out, _ = chunk_log_linear_attn(q, k, v, g, level_scales, cu_seqlens=cu_seqlens) + + o = [] + for i in range(cu_seqlens.shape[0] - 1): + bos, eos = cu_seqlens[i], cu_seqlens[i + 1] + v_s = v[:, bos:eos] + g_s = g[:, bos:eos] + k_s = k[:, bos:eos] + q_s = q[:, bos:eos] + level_scales_s = level_scales[:, bos:eos] + + o.append(naive_log_linear_attn(q_s, k_s, v_s, g_s, level_scales_s)) + ref = torch.cat(o, dim=1) + + assert_close("o", ref, out, 0.004) diff --git a/code/flash-linear-attention/tests/ops/test_mesa.py b/code/flash-linear-attention/tests/ops/test_mesa.py new file mode 100644 index 0000000000000000000000000000000000000000..915ef296b1e1671f6aeb5e4a164fd75826883e3a --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_mesa.py @@ -0,0 +1,276 @@ + +import os + +import pytest +import torch +import torch.nn.functional as F + +from fla.ops.mesa_net import chunk_mesa_net, mesa_net_decoding_one_step, naive_mesa_net_decoding_one_step, naive_mesa_net_exact +from fla.utils import assert_close, device, device_platform, is_intel_alchemist + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'gate_range', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-gate_range{}-{}".format(*test)) + for test in [ + (1, 63, 1, 64, [0.8, 0.99], torch.float16), + (2, 500, 4, 60, [0.8, 0.99], torch.float16), + (2, 1024, 8, 128, [0.8, 0.99], torch.float16), + (2, 1024, 8, 128, [0.01, 0.1], torch.float16), + (2, 1024, 8, 128, [1, 1], torch.float16), + (4, 2048, 8, 64, [0.8, 0.99], torch.float16), + ] + ], +) +@pytest.mark.skipif( + device_platform == 'intel', + reason='Intel Triton Failure', +) +def test_chunk( + B: int, + T: int, + H: int, + D: int, + gate_range: tuple[float, float], + dtype: torch.dtype, +): + torch.manual_seed(42) + q = torch.rand(B, T, H, D, dtype=dtype) / 10 + k = F.normalize(torch.rand(B, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype) + v = torch.rand(B, T, H, D, dtype=dtype) / 10 + beta = torch.rand(B, T, H, dtype=dtype).sigmoid() + lower_gate, upper_gate = gate_range + g = torch.rand(B, T, H, dtype=dtype).float().uniform_(lower_gate, upper_gate).log() + lamb = torch.rand(H, D, dtype=dtype).sigmoid() * 0.75 + 0.25 + q, k, v, beta, g, lamb = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, beta, g, lamb)) + do = torch.rand_like(v) + + k_init_rand = torch.nn.functional.normalize(torch.rand(B, H, D, device=device, dtype=dtype), dim=-1, p=2) + h_kk_init = (k_init_rand.unsqueeze(-1) * k_init_rand.unsqueeze(-2)).detach().clone().float().requires_grad_(True) + h_kv_init = torch.rand(B, H, D, D, dtype=torch.float32, device=device).requires_grad_(True) + d_h_kk_final = torch.rand_like(h_kk_init) + d_h_kv_final = torch.rand_like(h_kv_init) + + tri, tri_kk_final, tri_kv_final = chunk_mesa_net( + q=q.clone(), + k=k.clone(), + v=v.clone(), + beta=beta.clone(), + g=g.clone(), + lamb=lamb.clone(), + max_CG_iteration=D, + h_kk_init=h_kk_init.clone(), + h_kv_init=h_kv_init.clone(), + output_final_state=True, + ) + + ((tri * do).sum() + (tri_kk_final * d_h_kk_final).sum() + (tri_kv_final * d_h_kv_final).sum()).backward(retain_graph=True) + tri_dq, tri_dk, tri_dv, tri_dbeta, tri_dg, tri_dlamb = q.grad, k.grad, v.grad, beta.grad, g.grad, lamb.grad + tri_dh_kk_init, tri_dh_kv_init = h_kk_init.grad, h_kv_init.grad + q.grad = k.grad = v.grad = beta.grad = g.grad = lamb.grad = h_kk_init.grad = h_kv_init.grad = None + + ref, ref_hkk_final, ref_hkv_final = naive_mesa_net_exact( + q=q.clone(), + k=k.clone(), + v=v.clone(), + beta=beta.clone(), + g=g.clone(), + lamb=lamb.clone(), + h_kk_init=h_kk_init.clone(), + h_kv_init=h_kv_init.clone(), + ) + + ((ref * do).sum() + + (ref_hkk_final * d_h_kk_final).sum() + (ref_hkv_final * d_h_kv_final).sum()).backward(retain_graph=True) + ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dg, ref_dlamb = q.grad, k.grad, v.grad, beta.grad, g.grad, lamb.grad + ref_dh_kk_init, ref_dh_kv_init = h_kk_init.grad, h_kv_init.grad + q.grad = k.grad = v.grad = beta.grad = g.grad = lamb.grad = h_kk_init.grad = h_kv_init.grad = None + + assert_close('o', ref, tri, 0.006) + assert_close('h_kk_final', ref_hkk_final, tri_kk_final, 0.008) + assert_close('h_kv_final', ref_hkv_final, tri_kv_final, 0.008) + assert_close('dq', ref_dq, tri_dq, 0.008) + assert_close('dk', ref_dk, tri_dk, 0.008) + assert_close('dv', ref_dv, tri_dv, 0.008) + assert_close('db', ref_dbeta, tri_dbeta, 0.008) + assert_close('dg', ref_dg, tri_dg, 0.008) + assert_close('dlamb', ref_dlamb, tri_dlamb, 0.015) + assert_close('dh_kk_init', ref_dh_kk_init, tri_dh_kk_init, 0.008) + assert_close('dh_kv_init', ref_dh_kv_init, tri_dh_kv_init, 0.008) + + +@pytest.mark.parametrize( + ('H', 'D', 'gate_range', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-D{}-gate_range{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (3, 50, [0.8, 0.99], [0, 15], torch.float16), + (4, 64, [0.8, 0.99], [0, 14, 121, 421, 500], torch.float16), + (4, 64, [0.01, 0.1], [0, 256, 500, 1000], torch.float16), + (4, 100, [1, 1], [0, 15, 100, 300, 1200, 2000], torch.float16), + ] + ], +) +@pytest.mark.skipif( + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set', +) +def test_chunk_varlen( + H: int, + D: int, + gate_range: tuple[float, float], + cu_seqlens: list[int], + dtype: torch.dtype, +): + if is_intel_alchemist and D > 128: + pytest.skip(reason='chunk_gated_delta_rule is not supported on alchemist for D>128') + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + # randomly split the sequence into N segments + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.long, device=device) + T = cu_seqlens[-1] + N = len(cu_seqlens) - 1 + # seq-first required for inputs with variable lengths + q = torch.randn((1, T, H, D), dtype=dtype) / 10 + k = F.normalize(torch.randn(1, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype) + v = torch.randn((1, T, H, D), dtype=dtype) / 10 + lower_gate, upper_gate = gate_range + g = torch.rand(1, T, H, dtype=dtype).float().uniform_(lower_gate, upper_gate).log() + beta = torch.rand(1, T, H, dtype=dtype).sigmoid() + lamb = torch.rand(H, D, dtype=dtype).sigmoid() * 0.75 + 0.25 + + k_init_rand = torch.nn.functional.normalize(torch.rand(N, H, D, device=device, dtype=dtype), dim=-1, p=2) + h_kk_init = (k_init_rand.unsqueeze(-1) * k_init_rand.unsqueeze(-2)).detach().clone().float().requires_grad_(True) + h_kv_init = torch.rand(N, H, D, D, dtype=torch.float32, device=device).requires_grad_(True) + + q, k, v, beta, g, lamb, h_kk_init, h_kv_init = map(lambda x: x.to( + device).requires_grad_(), (q, k, v, beta, g, lamb, h_kk_init, h_kv_init)) + do = torch.rand_like(v) / 10 + d_h_kk_final = torch.rand_like(h_kk_init) + d_h_kv_final = torch.rand_like(h_kv_init) + + tri, tri_h_kk_final, tri_h_kv_final = chunk_mesa_net( + q=q.clone(), + k=k.clone(), + v=v.clone(), + beta=beta.clone(), + g=g.clone(), + lamb=lamb.clone(), + h_kk_init=h_kk_init.clone(), + h_kv_init=h_kv_init.clone(), + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + + ((tri * do).sum() + + (tri_h_kk_final * d_h_kk_final).sum() + (tri_h_kv_final * d_h_kv_final).sum()).backward(retain_graph=True) + tri_dq, tri_dk, tri_dv, tri_dbeta, tri_dg, tri_dlamb, tri_dh_kk_init, tri_dh_kv_init = \ + q.grad, k.grad, v.grad, beta.grad, g.grad, lamb.grad, h_kk_init.grad, h_kv_init.grad + q.grad = k.grad = v.grad = beta.grad = g.grad = lamb.grad = h_kk_init.grad = h_kv_init.grad = None + + ref = [] + ref_h_kk_t = [] + ref_h_kv_t = [] + for i in range(N): + ref_i, ref_h_kk_i, ref_h_kv_i = naive_mesa_net_exact( + q=q[:, cu_seqlens[i]:cu_seqlens[i+1]], + k=k[:, cu_seqlens[i]:cu_seqlens[i+1]], + v=v[:, cu_seqlens[i]:cu_seqlens[i+1]], + beta=beta[:, cu_seqlens[i]:cu_seqlens[i+1]], + g=g[:, cu_seqlens[i]:cu_seqlens[i+1]], + lamb=lamb, + h_kk_init=h_kk_init[i], + h_kv_init=h_kv_init[i], + ) + ref.append(ref_i) + ref_h_kk_t.append(ref_h_kk_i) + ref_h_kv_t.append(ref_h_kv_i) + ref = torch.cat(ref, 1) + ref_h_kk_t = torch.cat(ref_h_kk_t, 0) + ref_h_kv_t = torch.cat(ref_h_kv_t, 0) + + ((ref * do).sum() + (ref_h_kk_t * d_h_kk_final).sum() + (ref_h_kv_t * d_h_kv_final).sum()).backward(retain_graph=True) + ref_dq, ref_dk, ref_dv, ref_dbeta, ref_dg, ref_dlamb, ref_dh_kk_init, ref_dh_kv_init = \ + q.grad, k.grad, v.grad, beta.grad, g.grad, lamb.grad, h_kk_init.grad, h_kv_init.grad + q.grad = k.grad = v.grad = beta.grad = g.grad = lamb.grad = h_kk_init.grad = h_kv_init.grad = None + + assert_close('o', ref, tri, 0.006) + assert_close('h_kk_final', ref_h_kk_t, tri_h_kk_final, 0.008) + assert_close('h_kv_final', ref_h_kv_t, tri_h_kv_final, 0.008) + assert_close('dq', ref_dq, tri_dq, 0.008) + assert_close('dk', ref_dk, tri_dk, 0.008) + assert_close('dv', ref_dv, tri_dv, 0.008) + assert_close('db', ref_dbeta, tri_dbeta, 0.015) + assert_close('dlamb', ref_dlamb, tri_dlamb, 0.015) + assert_close('dg', ref_dg, tri_dg, 0.015) + assert_close('dh_kk_0', ref_dh_kk_init, tri_dh_kk_init, 0.007) + assert_close('dh_kv_0', ref_dh_kv_init, tri_dh_kv_init, 0.007) + + +@pytest.mark.parametrize( + ('B', 'H', 'D', 'gate_range', 'max_CG_step', 'dtype'), + [ + pytest.param(*test, id="B{}-H{}-D{}-gate_range{}-max_CG_step{}-{}".format(*test)) + for test in [ + (1, 3, 50, [0.95, 0.99], 1, torch.float16), + (2, 4, 60, [0.95, 0.99], 5, torch.float16), + (2, 8, 128, [0.95, 0.99], 1, torch.float16), + (2, 8, 128, [0.95, 0.99], 5, torch.float16), + (2, 8, 128, [0.95, 0.99], 30, torch.float16), + ] + ], +) +def test_decoding_one_step( + B: int, + H: int, + D: int, + gate_range: tuple[float, float], + max_CG_step: int, + dtype: torch.dtype, +): + if is_intel_alchemist and D > 128: + pytest.skip(reason='chunk_gated_delta_rule is not supported on alchemist for D>128') + torch.manual_seed(42) + torch.set_default_device(device) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + # randomly split the sequence into N segments + q = torch.rand((B, H, D), dtype=dtype) + k = F.normalize(torch.randn(B, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype) + v = torch.rand((B, H, D), dtype=dtype) + lower_gate, upper_gate = gate_range + g = torch.rand(B, H, dtype=dtype).float().uniform_(lower_gate, upper_gate).log() + beta = torch.rand(B, H, dtype=dtype).sigmoid() + lamb = torch.rand(H, D, dtype=dtype).sigmoid() * 0.75 + 0.25 + + k_init_rand = torch.nn.functional.normalize(torch.rand(B, H, D, device=device, dtype=dtype), dim=-1, p=2) + prev_h_kk = (k_init_rand.unsqueeze(-1) * k_init_rand.unsqueeze(-2)).detach().clone().float().requires_grad_(True) + prev_h_kv = torch.rand(B, H, D, D, dtype=torch.float32, device=device).requires_grad_(True) + + o, curr_h_kk, curr_h_kv = mesa_net_decoding_one_step( + q=q.clone(), + k=k.clone(), + v=v.clone(), + g=g.clone(), + lamb=lamb.clone(), + beta=beta.clone(), + prev_h_kk=prev_h_kk.clone(), + prev_h_kv=prev_h_kv.clone(), + max_CG_iteration=max_CG_step, + ) + + o_ref, curr_h_kk_re, curr_h_kv_re = naive_mesa_net_decoding_one_step( + q=q.clone(), + k=k.clone(), + v=v.clone(), + g=g.clone(), + lamb=lamb.clone(), + beta=beta.clone(), + prev_h_kk=prev_h_kk.clone(), + prev_h_kv=prev_h_kv.clone(), + max_CG_iteration=max_CG_step, + ) + + assert_close('o', o, o_ref, 0.005) + assert_close('curr_h_kk', curr_h_kk, curr_h_kk_re, 0.005) + assert_close('curr_h_kv', curr_h_kv, curr_h_kv_re, 0.005) diff --git a/code/flash-linear-attention/tests/ops/test_nsa.py b/code/flash-linear-attention/tests/ops/test_nsa.py new file mode 100644 index 0000000000000000000000000000000000000000..feaacdd3f44d9179238f4a5487206efebcfdf3b7 --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_nsa.py @@ -0,0 +1,148 @@ + +import os + +import pytest +import torch +import triton + +from fla.ops.nsa.naive import naive_nsa +from fla.ops.nsa.parallel import parallel_nsa +from fla.ops.utils import prepare_token_indices +from fla.utils import assert_close, device + + +# FIXME +@pytest.mark.parametrize( + ('B', 'T', 'H', 'HQ', 'D', 'S', 'block_size', 'scale', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-HQ{}-D{}-S{}-block_size{}-scale{}-{}".format(*test)) + for test in [ + (1, 63, 1, 16, 64, 16, 32, 1.0, torch.float16), + (3, 111, 1, 32, 100, 16, 32, 1.0, torch.float16), + (3, 1024, 2, 32, 60, 16, 32, 0.1, torch.float16), + (3, 1024, 2, 32, 128, 16, 32, 0.1, torch.float16), + (4, 2048, 2, 32, 64, 16, 32, 0.1, torch.float16), + ] + ], +) +def test_parallel( + B: int, + T: int, + H: int, + HQ: int, + D: int, + S: int, + block_size: int, + scale: float, + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + + q = torch.randn((B, T, HQ, D), dtype=dtype, device=device).requires_grad_(True) + k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True) + v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True) + do = torch.randn((B, T, HQ, D), dtype=dtype, device=device) + + block_indices = torch.full((B, T, H, S), T, dtype=torch.long, device=device) + for b in range(B): + for t in range(T): + for h in range(H): + i_i = torch.randperm(max(1, triton.cdiv(t, block_size)))[:S] + block_indices[b, t, h, :len(i_i)] = i_i + block_indices = block_indices.sort(-1)[0] + + ref = naive_nsa(q=q, k=k, v=v, block_indices=block_indices, block_size=block_size, scale=scale) + ref.backward(do) + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + + tri = parallel_nsa(q=q, k=k, v=v, block_indices=block_indices, block_size=block_size, scale=scale) + tri.backward(do) + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + + assert_close(" o", ref, tri, 0.005) + assert_close("dq", ref_dq, tri_dq, 0.005) + assert_close("dk", ref_dk, tri_dk, 0.005) + assert_close("dv", ref_dv, tri_dv, 0.005) + + +@pytest.mark.parametrize( + ('H', 'HQ', 'D', 'S', 'block_size', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-HQ{}-D{}-S{}-block_size{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (1, 16, 64, 16, 32, [0, 15], torch.float16), + (2, 32, 64, 16, 32, [0, 256, 500, 1000], torch.float16), + (2, 32, 100, 16, 32, [0, 15, 100, 300, 1200, 2000], torch.float16), + ] + ], +) +@pytest.mark.skipif( + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test because SKIP_TEST_CHUNK_VARLEN is set', +) +def test_parallel_varlen( + H: int, + HQ: int, + D: int, + S: int, + block_size: int, + cu_seqlens: list[int], + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + + T = cu_seqlens[-1] + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + + # seq-first required for inputs with variable lengths + q = torch.randn((1, T, HQ, D), dtype=dtype, device=device).requires_grad_() + k = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + do = torch.randn((1, T, HQ, D), dtype=dtype, device=device) + + block_indices = torch.full((1, T, H, S), T, dtype=torch.long, device=device) + seq_indices = prepare_token_indices(cu_seqlens).tolist() + + for i in range(T): + _, t = seq_indices[i] + for h in range(H): + i_i = torch.randperm(max(1, triton.cdiv(t, block_size)))[:S] + block_indices[0, i, h, :len(i_i)] = i_i + block_indices = block_indices.sort(-1)[0] + + ref = naive_nsa( + q=q, + k=k, + v=v, + block_indices=block_indices, + block_size=block_size, + cu_seqlens=cu_seqlens, + ) + ref.backward(do) + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + + tri = parallel_nsa( + q=q, + k=k, + v=v, + block_indices=block_indices, + block_size=block_size, + cu_seqlens=cu_seqlens, + ) + tri.backward(do) + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + + assert_close('o', ref, tri, 0.004) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) diff --git a/code/flash-linear-attention/tests/ops/test_path_attn.py b/code/flash-linear-attention/tests/ops/test_path_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..201430fe7656498e7a7237ea8fb14ea7cf132aa4 --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_path_attn.py @@ -0,0 +1,211 @@ + +import os + +import pytest +import torch +import torch.nn.functional as F +from einops import rearrange + +from fla.ops.path_attn.parallel import parallel_path_attention +from fla.utils import assert_close, device, is_intel_alchemist + + +def naive_path_attn(q, k, v, w, beta, g, scale, BT=64): + original_dtype = q.dtype + HQ = q.shape[2] + H = k.shape[2] + q, k, v, w, beta, g = map(lambda x: x.to(torch.float).transpose(1, 2), [q, k, v, w, beta, g]) + g_cumsum = g.cumsum(-1) + q = q.unsqueeze(2).expand(-1, -1, HQ//HQ, -1, -1).flatten(1, 2) + k = k.unsqueeze(2).expand(-1, -1, HQ//H, -1, -1).flatten(1, 2) + v = v.unsqueeze(2).expand(-1, -1, HQ//H, -1, -1).flatten(1, 2) + w = w.unsqueeze(2).expand(-1, -1, HQ//H, -1, -1).flatten(1, 2) + beta = beta.unsqueeze(2).expand(-1, -1, HQ//H, -1).flatten(1, 2) + + b, h, l, _ = q.shape + if l % BT != 0: + padding_size = BT - l % BT + q, k, w = map(lambda x: F.pad(x, (0, 0, 0, padding_size)), [q, k, w]) + beta = F.pad(beta, (0, padding_size)) + seq_len = q.shape[2] + w_beta = w * beta[..., None] + q, k, w, w_beta = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=BT), [q, k, w, w_beta]) + mask = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=0) + T = -(w_beta @ w.transpose(-1, -2)).masked_fill(mask, 0) + for i in range(1, BT): + T[..., i, :i] = T[..., i, :i].clone() + (T[..., i, :, None].clone() * T[..., :, :i].clone()).sum(-2) + T = T + torch.eye(BT, dtype=q.dtype, device=q.device) + Twbk = T @ (w_beta @ k.transpose(-1, -2)).masked_fill(mask, 0) + qw = (q @ w.transpose(-1, -2)).tril() + Twb = T @ w_beta + A_local = (q @ k.transpose(-1, -2)).tril() - qw @ Twbk + q = q - qw @ Twb + k = k - Twbk.transpose(-1, -2) @ w + H = w.transpose(-1, -2) @ Twb + A = torch.zeros(b, h, seq_len, seq_len, device=q.device) + q, k, w, w_beta = map(lambda x: rearrange(x, 'b h n c d -> b h (n c) d'), [q, k, w, w_beta]) + for i in range(0, seq_len, BT): + q_i = q[:, :, i:i+BT].clone() + for j in range(i - BT, -BT, -BT): + k_j = k[:, :, j:j+BT] + A_ij = q_i @ k_j.transpose(-1, -2) + A[:, :, i:i+BT, j:j+BT] = A_ij + q_i = q_i - q_i @ H[:, :, j // BT] + for i in range(0, seq_len//BT): + A[:, :, i*BT:i*BT+BT, i*BT:i*BT+BT] = A_local[:, :, i] + A = A.masked_fill_(~torch.tril(torch.ones(seq_len, seq_len, device=q.device, dtype=torch.bool)), float("-inf")) + A = A[:, :, :l, :l] + A = A + g_cumsum[..., None] - g_cumsum[..., None, :] + ref_o = (A * scale).softmax(-1).to(v) @ v + return ref_o.to(original_dtype).transpose(1, 2) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'HQ', 'D', 'use_forget_gate', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-HQ{}-D{}-use_forget_gate{}-{}".format(*test)) + for test in [ + # SY (2025/07/08): It somehow failed on Hopper with error msg: Aborted (core dumped) + # (10, 62, 2, 8, 128, True, torch.bfloat16), + (5, 512, 2, 8, 128, True, torch.bfloat16), + (3, 1024, 2, 8, 64, True, torch.bfloat16), + (2, 2000, 1, 4, 64, False, torch.bfloat16), + (1, 4000, 1, 2, 128, False, torch.bfloat16), + ] + ], +) +@pytest.mark.skipif( + is_intel_alchemist, + reason="Intel Triton Failure", +) +def test_parallel( + B: int, + H: int, + HQ: int, + T: int, + D: int, + use_forget_gate: bool, + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + + q = torch.randn((B, T, HQ, D), dtype=dtype, device=device).requires_grad_(True) + k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True) + v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True) + w = F.normalize(torch.randn((B, T, H, D), dtype=torch.float, device=device), dim=-1, p=2).requires_grad_(True) + beta = torch.empty((B, T, H), dtype=torch.float, device=device).uniform_(1.5, 2.0).requires_grad_(True) + if use_forget_gate: + g = torch.empty((B, T, HQ), dtype=torch.float, device=device).uniform_( + 0.95, 1).log().requires_grad_(True) + else: + g = None + do = torch.rand((B, T, HQ, D), dtype=dtype, device=device) + scale = D ** -0.5 + ref = naive_path_attn(q, k, v, w, beta, torch.zeros(B, T, HQ, device=device, dtype=torch.float) if g is None else g, scale) + ref.backward(do) + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + if use_forget_gate: + ref_dg, g.grad = g.grad.clone(), None + ref_dw, w.grad = w.grad.clone(), None + ref_db, beta.grad = beta.grad.clone(), None + + tri, _ = parallel_path_attention(q=q, k=k, v=v, w=w, beta=beta, g=g, scale=scale) + tri.backward(do) + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + if use_forget_gate: + tri_dg, g.grad = g.grad.clone(), None + tri_dw, w.grad = w.grad.clone(), None + tri_db, beta.grad = beta.grad.clone(), None + + assert_close(" o", ref, tri, 0.005) + assert_close("dq", ref_dq, tri_dq, 0.008) + assert_close("dk", ref_dk, tri_dk, 0.008) + assert_close("dv", ref_dv, tri_dv, 0.008) + if use_forget_gate: + assert_close("dg", ref_dg, tri_dg, 0.02) + assert_close("dw", ref_dw, tri_dw, 0.015) + assert_close("db", ref_db, tri_db, 0.02) + + +@pytest.mark.parametrize( + ('H', 'HQ', 'D', 'use_forget_gate', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-HQ{}-D{}-use_forget_gate{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (2, 4, 128, False, [0, 15, 333, 2048], torch.float16), + (2, 4, 128, True, [0, 15, 333, 2048], torch.float16), + (2, 4, 64, True, [0, 841, 889, 4096], torch.float16), + (2, 4, 64, False, [0, 841, 889, 2000, 3000, 4096], torch.float16), + (2, 16, 128, True, [0, 500, 1023, 2000, 3000, 4096], torch.float16), + ] + ], +) +@pytest.mark.skipif( + os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", + reason="Skipping test because TEST_CHUNK_VARLEN is enabled", +) +@pytest.mark.skipif( + is_intel_alchemist, + reason="Intel Triton Failure", +) +def test_parallel_varlen( + H: int, + HQ: int, + D: int, + use_forget_gate: bool, + cu_seqlens: list[int], + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + T = cu_seqlens[-1] + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + + q = torch.randn((1, T, HQ, D), dtype=dtype, device=device).requires_grad_(True) + k = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_(True) + v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_(True) + w = F.normalize(torch.randn((1, T, H, D), dtype=torch.float, device=device), dim=-1, p=2).requires_grad_(True) + beta = torch.rand((1, T, H), dtype=torch.float, device=device).sigmoid().requires_grad_(True) + if use_forget_gate: + g = torch.empty((1, T, HQ), dtype=torch.float, device=device).uniform_(0.95, 1).log().requires_grad_(True) + else: + g = None + do = torch.randn((1, T, HQ, D), dtype=dtype, device=device) + scale = D ** -0.5 + ref = torch.zeros(1, T, HQ, D, device=device, dtype=dtype) + for bos, eos in zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False): + g_segment = torch.zeros(1, eos - bos, HQ, device=device, dtype=torch.float) if g is None else g[:, bos:eos] + ref[:, bos:eos] = naive_path_attn( + q[:, bos:eos], k[:, bos:eos], v[:, bos:eos], + w[:, bos:eos], beta[:, bos:eos], g_segment, scale, + ) + ref.backward(do) + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + if use_forget_gate: + ref_dg, g.grad = g.grad.clone(), None + ref_dw, w.grad = w.grad.clone(), None + ref_db, beta.grad = beta.grad.clone(), None + tri, _ = parallel_path_attention(q=q, k=k, v=v, w=w, beta=beta, g=g, scale=scale, cu_seqlens=cu_seqlens) + tri.backward(do) + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + if use_forget_gate: + tri_dg, g.grad = g.grad.clone(), None + tri_dw, w.grad = w.grad.clone(), None + tri_db, beta.grad = beta.grad.clone(), None + assert_close(" o", ref, tri, 0.005) + assert_close("dq", ref_dq, tri_dq, 0.005) + assert_close("dk", ref_dk, tri_dk, 0.005) + assert_close("dv", ref_dv, tri_dv, 0.005) + if use_forget_gate: + assert_close("dg", ref_dg, tri_dg, 0.005) + assert_close("dw", ref_dw, tri_dw, 0.005) + assert_close("db", ref_db, tri_db, 0.005) diff --git a/code/flash-linear-attention/tests/ops/test_retention.py b/code/flash-linear-attention/tests/ops/test_retention.py new file mode 100644 index 0000000000000000000000000000000000000000..1313f47322c874fdf1ca2ddc79fe2c5a677e676a --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_retention.py @@ -0,0 +1,308 @@ + +import os + +import pytest +import torch + +from fla.ops.retention import chunk_retention, fused_chunk_retention, fused_recurrent_retention, parallel_retention +from fla.utils import assert_close, device + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'K', 'expand_ratio', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-K{}-expand_ratio{}-{}".format(*test)) + for test in [ + (1, 63, 1, 64, 1, torch.float16), + (2, 500, 3, 60, 1, torch.float16), + (2, 1000, 3, 100, 1, torch.float16), + (2, 1000, 3, 128, 2, torch.float16), + (3, 1024, 4, 256, 2, torch.float16), + (4, 2048, 4, 64, 2, torch.float16), + ] + ], +) +def test_chunk( + B: int, + T: int, + H: int, + K: int, + expand_ratio: int, + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + V = K * expand_ratio + + q = torch.randn((B, T, H, K), dtype=dtype, device=device).requires_grad_() + k = torch.randn((B, T, H, K), dtype=dtype, device=device).requires_grad_() + v = torch.randn((B, T, H, V), dtype=dtype, device=device).requires_grad_() + h0 = torch.randn((B, H, K, V), dtype=dtype, device=device).requires_grad_() + + do = torch.randn_like(v) + dht = torch.randn_like(h0) + ref, ref_ht = fused_recurrent_retention(q, k, v, initial_state=h0, output_final_state=True) + ((ref * do).sum() + (ref_ht * dht).sum()).backward() + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + + tri, tri_ht = chunk_retention(q, k, v, initial_state=h0, output_final_state=True) + ((tri * do).sum() + (tri_ht * dht).sum()).backward() + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) + + +@pytest.mark.parametrize( + ('H', 'K', 'expand_ratio', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-K{}-expand_ratio{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (4, 64, 1, [0, 15], torch.float16), + (4, 64, 2, [0, 256, 500, 1000], torch.float16), + (4, 100, 2, [0, 15, 100, 300, 1200, 2000], torch.float16), + ] + ], +) +@pytest.mark.skipif( + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set', +) +def test_chunk_varlen( + H: int, + K: int, + expand_ratio: int, + cu_seqlens: list[int], + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + V = K * expand_ratio + + N = len(cu_seqlens) - 1 + T = cu_seqlens[-1] + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.long, device=device) + + # seq-first required for inputs with variable lengths + q = torch.randn((1, T, H, K), dtype=dtype, device=device).requires_grad_() + k = torch.randn((1, T, H, K), dtype=dtype, device=device).requires_grad_() + v = torch.randn((1, T, H, V), dtype=dtype, device=device).requires_grad_() + h0 = torch.randn((N, H, K, V), dtype=dtype, device=device).requires_grad_() + do = torch.randn_like(v) + dht = torch.randn_like(h0) + + ref, ref_ht = fused_recurrent_retention( + q=q, + k=k, + v=v, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + ((ref * do).sum() + (ref_ht * dht).sum()).backward() + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_dh0, h0.grad = h0.grad.clone(), None + + tri, tri_ht = chunk_retention( + q=q, + k=k, + v=v, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + ((tri * do).sum() + (tri_ht * dht).sum()).backward() + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_dh0, h0.grad = h0.grad.clone(), None + + assert_close('o', ref, tri, 0.004) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) + assert_close('dh0', ref_dh0, tri_dh0, 0.005) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'K', 'expand_ratio', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-K{}-expand_ratio{}-{}".format(*test)) + for test in [ + (1, 63, 1, 64, 1, torch.float16), + (2, 500, 3, 60, 1, torch.float16), + (2, 1000, 3, 100, 1, torch.float16), + (2, 1000, 3, 128, 2, torch.float16), + (3, 1024, 4, 256, 2, torch.float16), + (4, 2048, 4, 64, 2, torch.float16), + ] + ], +) +def test_fused_chunk( + B: int, + T: int, + H: int, + K: int, + expand_ratio: int, + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + V = K * expand_ratio + + q = torch.randn((B, T, H, K), dtype=dtype, device=device).requires_grad_() + k = torch.randn((B, T, H, K), dtype=dtype, device=device).requires_grad_() + v = torch.randn((B, T, H, V), dtype=dtype, device=device).requires_grad_() + h0 = torch.randn((B, H, K, V), dtype=dtype, device=device).requires_grad_() + + do = torch.randn_like(v) + dht = torch.randn_like(h0) + ref, ref_ht = fused_recurrent_retention(q, k, v, initial_state=h0, output_final_state=True) + ((ref * do).sum() + (ref_ht * dht).sum()).backward() + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + + tri, tri_ht = fused_chunk_retention(q, k, v, initial_state=h0, output_final_state=True) + ((tri * do).sum() + (tri_ht * dht).sum()).backward() + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) + + +@pytest.mark.parametrize( + ('H', 'K', 'expand_ratio', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-K{}-expand_ratio{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (4, 64, 1, [0, 15], torch.float16), + (4, 64, 2, [0, 256, 500, 1000], torch.float16), + (4, 100, 2, [0, 15, 100, 300, 1200, 2000], torch.float16), + ] + ], +) +@pytest.mark.skipif( + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set', +) +def test_fused_chunk_varlen( + H: int, + K: int, + expand_ratio: int, + cu_seqlens: list[int], + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + V = K * expand_ratio + + N = len(cu_seqlens) - 1 + T = cu_seqlens[-1] + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.long, device=device) + + # seq-first required for inputs with variable lengths + q = torch.randn((1, T, H, K), dtype=dtype, device=device).requires_grad_() + k = torch.randn((1, T, H, K), dtype=dtype, device=device).requires_grad_() + v = torch.randn((1, T, H, V), dtype=dtype, device=device).requires_grad_() + h0 = torch.randn((N, H, K, V), dtype=dtype, device=device).requires_grad_() + do = torch.randn_like(v) + dht = torch.randn_like(h0) + + ref, ref_ht = fused_recurrent_retention( + q=q, + k=k, + v=v, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + ((ref * do).sum() + (ref_ht * dht).sum()).backward() + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_dh0, h0.grad = h0.grad.clone(), None + + tri, tri_ht = fused_chunk_retention( + q=q, + k=k, + v=v, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + ((tri * do).sum() + (tri_ht * dht).sum()).backward() + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_dh0, h0.grad = h0.grad.clone(), None + + assert_close('o', ref, tri, 0.004) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) + assert_close('dh0', ref_dh0, tri_dh0, 0.005) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'K', 'expand_ratio', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-K{}-expand_ratio{}-{}".format(*test)) + for test in [ + (1, 63, 1, 64, 1, torch.float16), + (2, 500, 4, 60, 1, torch.float16), + (2, 1024, 8, 128, 1, torch.float16), + (3, 1024, 8, 128, 2, torch.float16), + (3, 1024, 8, 256, 2, torch.float16), + (4, 2048, 8, 64, 2, torch.float16), + ] + ], +) +def test_parallel( + B: int, + T: int, + H: int, + K: int, + expand_ratio: int, + dtype: torch.dtype, +): + torch.manual_seed(42) + V = K * expand_ratio + + q = torch.randn((B, T, H, K), dtype=dtype, device=device).requires_grad_() + k = torch.randn((B, T, H, K), dtype=dtype, device=device).requires_grad_() + v = torch.randn((B, T, H, V), dtype=dtype, device=device).requires_grad_() + do = torch.randn_like(v) + + ref, _ = fused_recurrent_retention(q, k, v) + ref.backward(do) + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + + tri, _ = parallel_retention(q, k, v) + tri.backward(do) + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + + assert_close('o', ref, tri, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) diff --git a/code/flash-linear-attention/tests/ops/test_rwkv6.py b/code/flash-linear-attention/tests/ops/test_rwkv6.py new file mode 100644 index 0000000000000000000000000000000000000000..3848c2d7fb4774aa9c74facde281e5b2b53e5096 --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_rwkv6.py @@ -0,0 +1,192 @@ + +import os + +import pytest +import torch +import torch.nn.functional as F + +from fla.ops.rwkv6 import chunk_rwkv6 +from fla.ops.rwkv6.fused_recurrent import fused_recurrent_rwkv6 +from fla.utils import assert_close, device, device_platform + + +@pytest.mark.skipif( + device_platform == 'intel', + reason="Intel Triton Failure", +) +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'gate_logit_normalizer', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-gate_logit_normalizer{}-{}".format(*test)) + for test in [ + (1, 15, 2, 60, 1.0, torch.float16), + (3, 60, 3, 64, 0.1, torch.float16), + (3, 64, 2, 64, 1, torch.float16), + (4, 500, 3, 256, 1, torch.float16), + (4, 1000, 4, 64, 10, torch.float16), + (4, 2048, 4, 64, 1, torch.float16), + (4, 2048, 4, 256, 1, torch.float16), + ] + ], +) +def test_chunk( + B: int, + T: int, + H: int, + D: int, + gate_logit_normalizer: float, + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + + q = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() + k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() + v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() + w = F.logsigmoid(torch.randn((B, T, H, D), dtype=dtype, device=device)) / gate_logit_normalizer + + u = torch.randn(H, D, dtype=dtype, device=device).requires_grad_(True) + h0 = torch.randn(B, H, D, D, dtype=dtype, device=device).requires_grad_() + w = w.requires_grad_() + do = torch.randn_like(v) + + ref, ref_ht = fused_recurrent_rwkv6( + q.clone(), + k.clone(), + v.clone(), + w.clone(), + u.clone(), + initial_state=h0.clone(), + output_final_state=True, + ) + ref, _ = fused_recurrent_rwkv6( + q.clone(), + k.clone(), + v.clone(), + w.clone(), + u.clone(), + initial_state=h0.clone(), + output_final_state=False, + ) + + ((ref * do).sum()).backward() + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_dw, w.grad = w.grad.clone(), None + ref_du, u.grad = u.grad.clone(), None + ref_dh0, h0.grad = h0.grad.clone(), None + + # triton implementation + tri, tri_ht = chunk_rwkv6( + q.clone(), + k.clone(), + v.clone(), + w.clone(), + u.clone(), + initial_state=h0.clone(), + output_final_state=True, + ) + ((tri * do).sum()).backward() + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_dw, w.grad = w.grad.clone(), None + tri_du, u.grad = u.grad.clone(), None + tri_dh0, h0.grad = h0.grad.clone(), None + + assert_close('o', ref, tri, 0.004) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) + assert_close('dw', ref_dw, tri_dw, 0.005) + assert_close('du', ref_du, tri_du, 0.005) + assert_close('dh0', ref_dh0, tri_dh0, 0.005) + + +@pytest.mark.parametrize( + ('H', 'D', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-D{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (4, 64, [0, 15], torch.float16), + (4, 64, [0, 256, 500, 1000], torch.float16), + (4, 100, [0, 15, 100, 300, 1200, 2000], torch.float16), + ] + ], +) +def test_chunk_varlen( + H: int, + D: int, + cu_seqlens: list[int], + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + N = len(cu_seqlens) - 1 + T = cu_seqlens[-1] + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + + # seq-first required for inputs with variable lengths + q = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + k = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + w = F.logsigmoid(torch.randn((1, T, H, D), dtype=dtype, device=device)).requires_grad_(True) + u = torch.randn(H, D, dtype=dtype, device=device).requires_grad_(True) + h0 = torch.randn((N, H, D, D), dtype=dtype, device=device).requires_grad_() + do = torch.randn_like(v) + + ref, ref_ht = fused_recurrent_rwkv6( + q.clone(), + k.clone(), + v.clone(), + w.clone(), + u.clone(), + initial_state=h0.clone(), + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + ref, _ = fused_recurrent_rwkv6( + q.clone(), + k.clone(), + v.clone(), + w.clone(), + u.clone(), + initial_state=h0.clone(), + output_final_state=False, + cu_seqlens=cu_seqlens, + ) + ref.backward(do) + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_dw, w.grad = w.grad.clone(), None + ref_du, u.grad = u.grad.clone(), None + ref_dh0, h0.grad = h0.grad.clone(), None + + tri, tri_ht = chunk_rwkv6( + q.clone(), + k.clone(), + v.clone(), + w.clone(), + u.clone(), + initial_state=h0.clone(), + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + tri.backward(do) + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_dw, w.grad = w.grad.clone(), None + tri_du, u.grad = u.grad.clone(), None + tri_dh0, h0.grad = h0.grad.clone(), None + assert_close('o', ref, tri, 0.004) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) + assert_close('dw', ref_dw, tri_dw, 0.005) + assert_close('du', ref_du, tri_du, 0.005) + assert_close('dh0', ref_dh0, tri_dh0, 0.005) diff --git a/code/flash-linear-attention/tests/ops/test_rwkv7.py b/code/flash-linear-attention/tests/ops/test_rwkv7.py new file mode 100644 index 0000000000000000000000000000000000000000..16be77ae9c459596e0ca5a671e3ab99897f2d98d --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_rwkv7.py @@ -0,0 +1,300 @@ + +import os + +import pytest +import torch +import torch.nn.functional as F + +from fla.ops.generalized_delta_rule.dplr.fused_recurrent import fused_recurrent_dplr_delta_rule +from fla.ops.rwkv7.channel_mixing import channel_mixing_rwkv7, channel_mixing_rwkv7_torch +from fla.ops.rwkv7.fused_addcmul import fused_addcmul_rwkv7, torch_addcmul_rwkv7 +from fla.ops.rwkv7.fused_k_update import fused_k_rwkv7, k_update_ref +from fla.ops.rwkv7.fused_recurrent import fused_mul_recurrent_rwkv7 +from fla.ops.rwkv7.gate_output_correction import gate_output_correction, gate_output_correction_ref +from fla.utils import assert_close, device, is_nvidia_hopper + + +@pytest.mark.parametrize("B", [2]) +@pytest.mark.parametrize("T", [1024]) +@pytest.mark.parametrize("n_embd", [1024]) +@pytest.mark.parametrize("dim_ffn", [4096]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("inplace", [True, False]) +@pytest.mark.parametrize("xprevdim", [2, 3]) +@pytest.mark.skipif( + os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", + reason="Skipping test because TEST_CHUNK_VARLEN is enabled", +) +def test_channel_mixing_gradients(B, T, n_embd, dim_ffn, dtype, inplace, xprevdim): + torch.manual_seed(42) + torch._dynamo.config.cache_size_limit = 512 + + x = torch.randn( + B, T, n_embd, device=device, dtype=dtype, requires_grad=True, + ) + if xprevdim == 3: + x_prev = torch.randn( + B, 1, n_embd, device=device, dtype=dtype, requires_grad=True, + ) + else: + x_prev = torch.randn( + B, n_embd, device=device, dtype=dtype, requires_grad=True, + ) + x_k = torch.randn(1, 1, n_embd, device=device, dtype=dtype, requires_grad=True) + K_ = torch.randn(n_embd, dim_ffn, device=device, dtype=dtype, requires_grad=True) + V_ = torch.randn(dim_ffn, n_embd, device=device, dtype=dtype, requires_grad=True) + + x2 = x.clone().detach().requires_grad_(True) + x_prev2 = x_prev.clone().detach().requires_grad_(True) + x_k2 = x_k.clone().detach().requires_grad_(True) + K_2 = K_.clone().detach().requires_grad_(True) + V_2 = V_.clone().detach().requires_grad_(True) + + o1, last1 = channel_mixing_rwkv7_torch( + x.to(torch.float32), + x_prev.to(torch.float32), + x_k.to(torch.float32), + K_.to(torch.float32), + V_.to(torch.float32), + ) + loss1 = o1.mean() + last1.mean() + loss1.backward() + + o2, last2 = channel_mixing_rwkv7(x2, x_prev2, x_k2, K_2, V_2, inplace) + loss2 = o2.mean() + last2.mean() + loss2.backward() + + assert_close(" dx", x.grad, x2.grad, ratio=5e-3) + assert_close(" dxprev", x_prev.grad, x_prev2.grad, ratio=5e-3) + assert_close(" dx_k", x_k.grad, x_k2.grad, ratio=5e-3) + assert_close(" dK_", K_.grad, K_2.grad, ratio=5e-3) + assert_close(" dV_", V_.grad, V_2.grad, ratio=5e-3) + + +@pytest.mark.parametrize('B', [2]) +@pytest.mark.parametrize('T', [1, 1024]) +@pytest.mark.parametrize('H', [1]) +@pytest.mark.parametrize('D', [64]) +@pytest.mark.parametrize('scale', [None, 1]) +@pytest.mark.parametrize('dtype', [torch.float32]) +@pytest.mark.skipif( + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '0', + reason='Skipping test because TEST_CHUNK_VARLEN is enabled', +) +def test_fused_mul_recurrent_fwd( + B: int, + T: int, + H: int, + D: int, + scale: float, + dtype: torch.dtype, +): + torch.manual_seed(42) + r = torch.empty(B, T, H, D, device=device).uniform_(-8, -6).to(dtype=dtype) + k = torch.empty(B, T, H, D, device=device).uniform_(-8, -6).to(dtype=dtype) + v = torch.empty(B, T, H, D, device=device).uniform_(-8, -6).to(dtype=dtype) + w = torch.empty(B, T, H, D, device=device).uniform_(-8, -6).to(dtype=dtype) + + kk = torch.empty(B, T, H, D, device=device).uniform_(-1, 1) + kk = F.normalize(kk, dim=-1).to(dtype=dtype) + + a = -kk.clone() + a_scale = torch.empty(B, T, H, D, device=device).uniform_(0, 0.1).to(dtype=dtype) + b = (kk * a_scale).requires_grad_(False) # kk*a + h0 = torch.randn(B, H, D, D, dtype=torch.float) + r, k, v, a, a_scale, b, w, h0 = map(lambda x: x.to(device).requires_grad_(False), + (r, k, v, a, a_scale, b, w, h0)) + ref, ref_ht = fused_recurrent_dplr_delta_rule( + q=r.clone(), + k=k.clone(), + v=v.clone(), + a=a.clone(), + b=b.clone(), + gk=w.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + + tri, tri_ht = fused_mul_recurrent_rwkv7( + r=r.clone(), + w=w.clone(), + k=k.clone(), + v=v.clone(), + kk=kk.clone(), + a=a_scale.clone(), + scale=scale, + initial_state=h0.clone(), + output_final_state=True, + ) + assert_close('o', ref, tri, 0.002) + assert_close('ht', ref_ht, tri_ht, 0.002) + + +@pytest.mark.parametrize("B", [1]) +@pytest.mark.parametrize("T", [20, 1024, 4100, 131072]) +@pytest.mark.parametrize("H", [2]) +@pytest.mark.parametrize("D", [64]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("use_g", [True, False]) +@pytest.mark.skipif( + os.getenv("SKIP_TEST_CHUNK_VARLEN") == "0", + reason="Skipping test because TEST_CHUNK_VARLEN is enabled", +) +def test_fused_rwkv7_addcmul( + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_g: bool, +): + if T == 128 * 1024 and not is_nvidia_hopper: + pytest.skip("Skipping test for T=131072 on non-Hopper GPUs") + hidden_size = H*D + hidden_states = torch.randn(B, T, hidden_size).to(device).to(dtype).requires_grad_() + xx = torch.randn(B, T, hidden_size).to(device).to(dtype).requires_grad_() + x_r = torch.randn(1, 1, hidden_size).to(device).to(dtype).requires_grad_() + x_w = torch.randn(1, 1, hidden_size).to(device).to(dtype).requires_grad_() + x_k = torch.randn(1, 1, hidden_size).to(device).to(dtype).requires_grad_() + x_v = torch.randn(1, 1, hidden_size).to(device).to(dtype).requires_grad_() + x_a = torch.randn(1, 1, hidden_size).to(device).to(dtype).requires_grad_() + if use_g: + x_g = torch.randn(1, 1, hidden_size).to(device).to(dtype).requires_grad_() + else: + x_g = None + xr0, xw0, xk0, xv0, xa0, xg0 = fused_addcmul_rwkv7(hidden_states, xx, x_r, x_w, x_k, x_v, x_a, x_g) + xr1, xw1, xk1, xv1, xa1, xg1 = torch_addcmul_rwkv7(hidden_states.float(), + xx.float(), x_r.float(), + x_w.float(), x_k.float(), + x_v.float(), x_a.float(), + x_g.float() if use_g else None) + ratio = 1e-5 if dtype == torch.float32 else 0.002 + assert_close("xr0", xr0, xr1, ratio=ratio) + assert_close("xw0", xw0, xw1, ratio=ratio) + assert_close("xk0", xk0, xk1, ratio=ratio) + assert_close("xv0", xv0, xv1, ratio=ratio) + assert_close("xa0", xa0, xa1, ratio=ratio) + if use_g: + assert_close("xg0", xg0, xg1, ratio=ratio) + (xr0 + xw0 + xk0 + xv0 + xa0 + xg0).sum().backward() + else: + (xr0 + xw0 + xk0 + xv0 + xa0).sum().backward() + d_ixr = x_r.grad.clone() + d_ixw = x_w.grad.clone() + d_ixk = x_k.grad.clone() + d_ixv = x_v.grad.clone() + d_ixa = x_a.grad.clone() + d_hidden = hidden_states.grad.clone() + d_xx = xx.grad.clone() + + x_r.grad.zero_() + x_w.grad.zero_() + x_k.grad.zero_() + x_v.grad.zero_() + x_a.grad.zero_() + if use_g: + d_ixg = x_g.grad.clone() + x_g.grad.zero_() + hidden_states.grad.zero_() + xx.grad.zero_() + + if use_g: + (xr1 + xw1 + xk1 + xv1 + xa1 + xg1).sum().backward() + else: + (xr1 + xw1 + xk1 + xv1 + xa1).sum().backward() + d_ixr1 = x_r.grad.clone() + d_ixw1 = x_w.grad.clone() + d_ixk1 = x_k.grad.clone() + d_ixv1 = x_v.grad.clone() + d_ixa1 = x_a.grad.clone() + if use_g: + d_ixg1 = x_g.grad.clone() + d_hidden1 = hidden_states.grad.clone() + d_xx1 = xx.grad.clone() + + assert_close("d_ixr", d_ixr, d_ixr1, ratio=ratio) + assert_close("d_ixw", d_ixw, d_ixw1, ratio=ratio) + assert_close("d_ixk", d_ixk, d_ixk1, ratio=ratio) + assert_close("d_ixv", d_ixv, d_ixv1, ratio=ratio) + assert_close("d_ixa", d_ixa, d_ixa1, ratio=ratio) + if use_g: + assert_close("d_ixg", d_ixg, d_ixg1, ratio=ratio) + assert_close("d_hidden", d_hidden, d_hidden1, ratio=ratio) + assert_close("d_xx", d_xx, d_xx1, ratio=ratio) + + +@pytest.mark.parametrize("B", [4]) +@pytest.mark.parametrize("T", [13, 4096, 8000]) +@pytest.mark.parametrize("H", [64]) +@pytest.mark.parametrize("D", [64]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("ka_shape", [1, 3]) +def test_fused_k_update( + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + ka_shape: int, +): + k = torch.randn(B, T, H*D).uniform_(-8, 8).to(device).to(dtype).requires_grad_() + a = torch.randn(B, T, H*D).uniform_(-8, 8).to(device).to(dtype).requires_grad_() + if ka_shape == 1: + ka = torch.randn(H*D).uniform_(-8, 8).to(device).to(dtype).requires_grad_() + else: + ka = torch.randn(1, 1, H*D).uniform_(-8, 8).to(device).to(dtype).requires_grad_() + + ref = k_update_ref(k.float(), a.float(), ka.float()) + ref.sum().backward() + ref_dk, k.grad = k.grad.clone(), None + ref_da, a.grad = a.grad.clone(), None + ref_dka, ka.grad = ka.grad.clone(), None + tri = fused_k_rwkv7(k, a, ka) + tri.sum().backward() + ratio = 5e-5 if dtype == torch.float32 else 0.002 + assert_close(" o", tri, ref, ratio=ratio) + assert_close(" dk", ref_dk, k.grad, ratio=ratio) + assert_close(" da", ref_da, a.grad, ratio=ratio) + assert_close("dka", ref_dka, ka.grad, ratio=ratio) + + +@pytest.mark.parametrize("B", [4]) +@pytest.mark.parametrize("T", [4096]) +@pytest.mark.parametrize("H", [64]) +@pytest.mark.parametrize("D", [64]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +def test_gate_output_correction( + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): + value_dim = H * D + torch.manual_seed(0) + + o_ref = torch.randn(B, T, value_dim, device=device, dtype=dtype, requires_grad=True) + r_ref = torch.randn(B, T, H, D, device=device, dtype=dtype, requires_grad=True) + k_ref = torch.randn(B, T, H, D, device=device, dtype=dtype, requires_grad=True) + r_k_ref = torch.randn(H, D, device=device, dtype=dtype, requires_grad=True) + v_ref = torch.randn(B, T, H, D, device=device, dtype=dtype, requires_grad=True) + g_ref = torch.randn(B, T, value_dim, device=device, dtype=dtype, requires_grad=True) + + tensors_cus = [t.clone().detach().requires_grad_(True) for t in [o_ref, r_ref, k_ref, r_k_ref, v_ref, g_ref]] + o_cus, r_cus, k_cus, r_k_cus, v_cus, g_cus = tensors_cus + + output_ref = gate_output_correction_ref(o_ref.float(), r_ref.float(), k_ref.float(), + r_k_ref.float(), v_ref.float(), g_ref.float()) + output_ref.sum().backward() + + output_cus = gate_output_correction(o_cus, r_cus, k_cus, r_k_cus, v_cus, g_cus) + output_cus.sum().backward() + + assert_close(" o", output_ref, output_cus, 0.002) + assert_close("do", o_ref.grad, o_cus.grad, 0.002) + assert_close("dr", r_ref.grad, r_cus.grad, 0.002) + assert_close("dk", k_ref.grad, k_cus.grad, 0.002) + assert_close("drk", r_k_ref.grad, r_k_cus.grad, 0.002) + assert_close("dv", v_ref.grad, v_cus.grad, 0.002) + assert_close("dg", g_ref.grad, g_cus.grad, 0.002) diff --git a/code/flash-linear-attention/tests/ops/test_simple_gla.py b/code/flash-linear-attention/tests/ops/test_simple_gla.py new file mode 100644 index 0000000000000000000000000000000000000000..b32792911355484e4a3b2ce4e0f0c1293a96542b --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_simple_gla.py @@ -0,0 +1,670 @@ + +import os + +import pytest +import torch +import torch.nn.functional as F + +from fla.ops.simple_gla.chunk import chunk_simple_gla +from fla.ops.simple_gla.fused_chunk import fused_chunk_simple_gla +from fla.ops.simple_gla.fused_recurrent import fused_recurrent_simple_gla +from fla.ops.simple_gla.naive import naive_parallel_simple_gla, naive_recurrent_simple_gla +from fla.ops.simple_gla.parallel import parallel_simple_gla +from fla.utils import assert_close, device + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-{}".format(*test)) + for test in [ + (1, 63, 1, 64, 1, 1, torch.float), + (2, 500, 4, 60, 1, 1, torch.float), + (2, 1024, 8, 128, 1, 0.1, torch.float), + (2, 1024, 8, 128, 0.1, 1, torch.float), + (2, 1024, 8, 128, 1, 10, torch.float), + (4, 2048, 8, 64, 0.1, 1, torch.float), + (2, 1024, 8, 128, 1, 0.1, torch.float16), + (2, 1024, 8, 128, 1, 10, torch.float16), + ] + ], +) +def test_fused_recurrent( + B: int, + T: int, + H: int, + D: int, + scale: float, + gate_logit_normalizer: float, + dtype: torch.dtype, +): + torch.manual_seed(42) + + q = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() + k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() + v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_() + g = torch.randn((B, T, H), dtype=dtype, device=device) + g = (F.logsigmoid(g) / gate_logit_normalizer).requires_grad_() + h0 = torch.randn(B, H, D, D, device=device).requires_grad_() + dht = torch.randn_like(h0) + do = torch.randn_like(v) + ref, ref_ht = naive_recurrent_simple_gla( + q=q, + k=k, + v=v, + g=g, + scale=scale, + initial_state=h0, + output_final_state=True, + ) + ((ref * do).sum() + (ref_ht * dht).sum()).backward() + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_dg, g.grad = g.grad.clone(), None + ref_dh0, h0.grad = h0.grad.clone(), None + + tri, tri_ht = fused_recurrent_simple_gla( + q=q, + k=k, + v=v, + g=g, + scale=scale, + initial_state=h0, + output_final_state=True, + ) + ((tri * do).sum() + (tri_ht * dht).sum()).backward() + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_dg, g.grad = g.grad.clone(), None + tri_dh0, h0.grad = h0.grad.clone(), None + + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) + assert_close('dg', ref_dg, tri_dg, 0.005, err_atol=2e-4) + assert_close('dh0', ref_dh0, tri_dh0, 0.005) + + +@pytest.mark.parametrize( + ('H', 'D', 'scale', 'gate_logit_normalizer', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-D{}-scale{}-gate_logit_normalizer{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (4, 64, 1, 1, [0, 15], torch.float), + (4, 64, 1, 1, [0, 256, 500, 1000], torch.float), + (4, 100, 0.1, 1, [0, 15, 100, 300, 1200, 2000], torch.float), + (4, 100, 1, 1, [0, 15, 100, 300, 1200, 2000], torch.float), + (4, 100, 1, 10, [0, 15, 100, 300, 1200, 2000], torch.float), + (4, 64, 1, 1, [0, 1, 100, 300, 1200, 2048], torch.float16), + (4, 128, 1, 1, [0, 200, 512, 1200, 2048], torch.float16), + ] + ], +) +def test_fused_recurrent_varlen( + H: int, + D: int, + scale: float, + gate_logit_normalizer: float, + cu_seqlens: list[int], + dtype: torch.dtype, +): + torch.manual_seed(42) + + N = len(cu_seqlens) - 1 + T = cu_seqlens[-1] + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + + q = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + k = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + g = torch.randn((1, T, H), dtype=dtype, device=device) + g = (F.logsigmoid(g) / gate_logit_normalizer).requires_grad_() + h0 = torch.randn(N, H, D, D, device=device).requires_grad_() + dht = torch.randn_like(h0) + do = torch.randn_like(v) + + refs, ref_hts = [], [] + for i, (bos, eos) in enumerate(zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False)): + ref, ref_ht = naive_recurrent_simple_gla( + q=q[:, bos:eos], + k=k[:, bos:eos], + v=v[:, bos:eos], + g=g[:, bos:eos], + scale=scale, + initial_state=h0[i], + output_final_state=True, + ) + refs.append(ref) + ref_hts.append(ref_ht) + ref = torch.cat(refs, 1) + ref_ht = torch.cat(ref_hts, 0) + ((ref * do).sum() + (ref_ht * dht).sum()).backward() + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_dg, g.grad = g.grad.clone(), None + ref_dh0, h0.grad = h0.grad.clone(), None + + tri, tri_ht = fused_recurrent_simple_gla( + q=q, + k=k, + v=v, + g=g, + scale=scale, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + ((tri * do).sum() + (tri_ht * dht).sum()).backward() + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_dg, g.grad = g.grad.clone(), None + tri_dh0, h0.grad = h0.grad.clone(), None + + assert_close('o', ref, tri, 0.005) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) + assert_close('dg', ref_dg, tri_dg, 0.005, err_atol=2e-4) + assert_close('dh0', ref_dh0, tri_dh0, 0.005) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-{}".format(*test)) + for test in [ + (1, 63, 1, 64, 1, 1, torch.float16), + (2, 500, 3, 60, 1, 1, torch.float16), + (1, 1000, 4, 128, 1, 0.1, torch.float16), + (2, 1000, 4, 128, 0.1, 1, torch.float16), + (3, 1000, 4, 128, 0.1, 10, torch.float16), + (4, 2048, 8, 64, 0.1, 1, torch.float16), + ] + ], +) +def test_chunk( + B: int, + T: int, + H: int, + D: int, + scale: float, + gate_logit_normalizer: float, + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + q = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True) + k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True) + v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True) + g = torch.randn((B, T, H), dtype=torch.float32, device=device) + h0 = torch.rand((B, H, D, D), dtype=torch.float32, device=device).requires_grad_(True) + dht = torch.randn_like(h0) + g = (F.logsigmoid(g) / gate_logit_normalizer).requires_grad_(True) + do = torch.randn_like(v) + + ref, ref_ht = fused_recurrent_simple_gla( + q=q, + k=k, + v=v, + g=g, + scale=scale, + initial_state=h0, + output_final_state=True, + ) + ((ref * do).sum() + (dht * ref_ht).sum()).backward() + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_dg, g.grad = g.grad.clone(), None + ref_dh0, h0.grad = h0.grad.clone(), None + + tri, tri_ht = chunk_simple_gla( + q=q, + k=k, + v=v, + g=g, + scale=scale, + initial_state=h0, + output_final_state=True, + ) + ((tri * do).sum() + (dht * tri_ht).sum()).backward() + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_dg, g.grad = g.grad.clone(), None + tri_dh0, h0.grad = h0.grad.clone(), None + + assert_close('o', ref, tri, 0.004) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) + assert_close('dg', ref_dg, tri_dg, 0.005) + assert_close('dh0', ref_dh0, tri_dh0, 0.005) + + +@pytest.mark.parametrize( + ('H', 'D', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-D{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (4, 64, [0, 15], torch.float16), + (4, 64, [0, 256, 500, 1000], torch.float16), + (4, 100, [0, 15, 100, 300, 1200, 2000], torch.float16), + ] + ], +) +@pytest.mark.skipif( + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set', +) +def test_chunk_varlen( + H: int, + D: int, + cu_seqlens: list[int], + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + + N = len(cu_seqlens) - 1 + T = cu_seqlens[-1] + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + + # seq-first required for inputs with variable lengths + q = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + k = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + g = F.logsigmoid(torch.randn((1, T, H), dtype=dtype, device=device)).requires_grad_() + h0 = torch.randn((N, H, D, D), dtype=torch.float32, device=device).requires_grad_() + dht = torch.randn_like(h0) + do = torch.randn_like(v) + + ref, ref_ht = fused_recurrent_simple_gla( + q=q, + k=k, + v=v, + g=g, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + ((ref * do).sum() + (dht * ref_ht).sum()).backward() + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_dg, g.grad = g.grad.clone(), None + ref_dh0, h0.grad = h0.grad.clone(), None + + tri, tri_ht = chunk_simple_gla( + q=q, + k=k, + v=v, + g=g, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + ((tri * do).sum() + (dht * tri_ht).sum()).backward() + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_dg, g.grad = g.grad.clone(), None + tri_dh0, h0.grad = h0.grad.clone(), None + + assert_close('o', ref, tri, 0.004) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) + assert_close('dg', ref_dg, tri_dg, 0.005) + assert_close('dh0', ref_dh0, tri_dh0, 0.005) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-{}".format(*test)) + for test in [ + (1, 63, 1, 64, 1, 1, torch.float16), + (2, 500, 3, 60, 1, 1, torch.float16), + (1, 1000, 4, 128, 1, 0.1, torch.float16), + (2, 1000, 4, 128, 0.1, 1, torch.float16), + (3, 1000, 4, 128, 0.1, 10, torch.float16), + (4, 2048, 8, 64, 0.1, 1, torch.float16), + ] + ], +) +def test_fused_chunk( + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + scale: float, + gate_logit_normalizer: float, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + q = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True) + k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True) + v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True) + g = torch.randn((B, T, H), dtype=torch.float32, device=device) + h0 = torch.rand((B, H, D, D), dtype=torch.float32, device=device).requires_grad_(True) + dht = torch.randn_like(h0) + g = (F.logsigmoid(g) / gate_logit_normalizer).requires_grad_(True) + do = torch.randn_like(v) + + ref, ref_ht = fused_recurrent_simple_gla( + q=q, + k=k, + v=v, + g=g, + scale=scale, + initial_state=h0, + output_final_state=True, + ) + ((ref * do).sum() + (dht * ref_ht).sum()).backward() + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_dg, g.grad = g.grad.clone(), None + ref_dh0, h0.grad = h0.grad.clone(), None + + tri, tri_ht = fused_chunk_simple_gla( + q=q, + k=k, + v=v, + g=g, + scale=scale, + initial_state=h0, + output_final_state=True, + ) + ((tri * do).sum() + (dht * tri_ht).sum()).backward() + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_dg, g.grad = g.grad.clone(), None + tri_dh0, h0.grad = h0.grad.clone(), None + + assert_close('o', ref, tri, 0.004) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) + assert_close('dg', ref_dg, tri_dg, 0.005) + assert_close('dh0', ref_dh0, tri_dh0, 0.005) + + +@pytest.mark.parametrize( + ('H', 'D', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-D{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (4, 64, [0, 15], torch.float16), + (4, 64, [0, 256, 500, 1000], torch.float16), + (4, 100, [0, 15, 100, 300, 1200, 2000], torch.float16), + ] + ], +) +@pytest.mark.skipif( + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set', +) +def test_fused_chunk_varlen( + H: int, + D: int, + cu_seqlens: list[int], + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + + N = len(cu_seqlens) - 1 + T = cu_seqlens[-1] + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + + q = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + k = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + g = F.logsigmoid(torch.randn((1, T, H), dtype=dtype, device=device)).requires_grad_() + h0 = torch.randn((N, H, D, D), dtype=torch.float32, device=device).requires_grad_() + dht = torch.randn_like(h0) + do = torch.randn_like(v) + + ref, ref_ht = fused_recurrent_simple_gla( + q=q, + k=k, + v=v, + g=g, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + ((ref * do).sum() + (dht * ref_ht).sum()).backward() + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_dg, g.grad = g.grad.clone(), None + ref_dh0, h0.grad = h0.grad.clone(), None + + tri, tri_ht = fused_chunk_simple_gla( + q=q, + k=k, + v=v, + g=g, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + ((tri * do).sum() + (dht * tri_ht).sum()).backward() + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_dg, g.grad = g.grad.clone(), None + tri_dh0, h0.grad = h0.grad.clone(), None + + assert_close('o', ref, tri, 0.004) + assert_close('ht', ref_ht, tri_ht, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) + assert_close('dg', ref_dg, tri_dg, 0.005) + assert_close('dh0', ref_dh0, tri_dh0, 0.005) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'scale', 'gate_logit_normalizer', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-gate_logit_normalizer{}-{}".format(*test)) + for test in [ + (1, 63, 1, 64, 1, 1, torch.float16), + (2, 500, 3, 60, 1, 1, torch.float16), + (2, 1024, 4, 128, 0.1, 1, torch.float16), + (3, 1024, 4, 128, 0.1, 10, torch.float16), + (3, 1024, 4, 256, 0.1, 0.1, torch.float16), + (4, 2048, 4, 64, 0.1, 0.1, torch.float16), + ] + ], +) +def test_parallel( + B: int, + T: int, + H: int, + D: int, + scale: float, + gate_logit_normalizer: float, + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + USE_G = gate_logit_normalizer > 0 + q = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True) + k = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True) + v = torch.randn((B, T, H, D), dtype=dtype, device=device).requires_grad_(True) + g = F.logsigmoid(torch.randn((B, T, H), dtype=dtype, device=device)) if USE_G else None + g = (g / gate_logit_normalizer).requires_grad_(True) if USE_G else None + do = torch.randn_like(v) + + ref, _ = fused_recurrent_simple_gla(q=q, k=k, v=v, g=g, scale=scale, output_final_state=True) + _, ref_A = naive_parallel_simple_gla(q=q, k=k, v=v, g=g, scale=scale) + ref.backward(do) + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + if USE_G: + ref_dg, g.grad = g.grad.clone(), None + + tri, tri_A = parallel_simple_gla(q=q, k=k, v=v, g=g, scale=scale, output_attentions=True) + tri.backward(do) + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + if USE_G: + tri_dg, g.grad = g.grad.clone(), None + assert_close('o', ref, tri, 0.005) + assert_close('A', ref_A, tri_A, 0.005) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) + if USE_G: + assert_close('dg', ref_dg, tri_dg, 0.015) + + +@pytest.mark.parametrize( + ('H', 'D', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-D{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (4, 64, [0, 15], torch.float16), + (4, 64, [0, 256, 500, 1000], torch.float16), + (4, 100, [0, 15, 100, 300, 1200, 2000], torch.float16), + ] + ], +) +@pytest.mark.skipif( + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set', +) +def test_parallel_varlen( + H: int, + D: int, + cu_seqlens: list[int], + dtype: torch.dtype, +): + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + + T = cu_seqlens[-1] + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + + q = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + k = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + v = torch.randn((1, T, H, D), dtype=dtype, device=device).requires_grad_() + g = F.logsigmoid(torch.randn((1, T, H), dtype=dtype, device=device)).requires_grad_() + do = torch.randn_like(v) + + ref, _ = fused_recurrent_simple_gla( + q=q, + k=k, + v=v, + g=g, + output_final_state=False, + cu_seqlens=cu_seqlens, + ) + ((ref * do).sum()).backward() + ref_dq, q.grad = q.grad.clone(), None + ref_dk, k.grad = k.grad.clone(), None + ref_dv, v.grad = v.grad.clone(), None + ref_dg, g.grad = g.grad.clone(), None + + tri, _ = parallel_simple_gla( + q=q, + k=k, + v=v, + g=g, + cu_seqlens=cu_seqlens, + ) + ((tri * do).sum()).backward() + tri_dq, q.grad = q.grad.clone(), None + tri_dk, k.grad = k.grad.clone(), None + tri_dv, v.grad = v.grad.clone(), None + tri_dg, g.grad = g.grad.clone(), None + + assert_close('o', ref, tri, 0.004) + assert_close('dq', ref_dq, tri_dq, 0.005) + assert_close('dk', ref_dk, tri_dk, 0.005) + assert_close('dv', ref_dv, tri_dv, 0.005) + assert_close('dg', ref_dg, tri_dg, 0.005) + + +@pytest.mark.parametrize( + ('vary_A', 'dtype'), + [ + pytest.param(True, torch.float, id=f'vary_A{True}-dtype{torch.float}'), + pytest.param(False, torch.float, id=f'vary_A{False}-dtype{torch.float}'), + pytest.param(True, torch.float16, id=f'vary_A{True}-dtype{torch.float16}'), + pytest.param(False, torch.float16, id=f'vary_A{False}-dtype{torch.float16}'), + ], +) +def test_simple_gla_to_mamba2(vary_A, dtype): + try: + from mamba_ssm.modules.ssd_minimal import ssd_minimal_discrete + from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined + except ImportError: + pytest.skip('mamba_ssm is not installed.') + torch.manual_seed(42) + + # Dimensions, Denoted (B, T, Q, D, P) in Mamba2 paper + batch, seq_len, chunk_size, dim, headdim = 2, 512, 8, 64, 16 + n_heads = dim // headdim # (H) in the paper + ngroups = n_heads # (G) in the paper; NOTE: do not use group-query here + dstate = 64 # (N) in the paper + atol = 5e-4 if dtype == torch.float else 1e-2 + + x = 0.1 * torch.randn(batch, seq_len, n_heads, headdim, dtype=dtype, device=device) + dt = torch.ones(batch, seq_len, n_heads, dtype=dtype, device=device) # dt=1 can be ignored + + if vary_A: + A = -0.1 * torch.rand(1, seq_len, n_heads, dtype=dtype, device=device) + else: # constant A for all position + A = -0.1 * torch.rand(n_heads, dtype=dtype, device=device) + + B = 0.1 * torch.randn(batch, seq_len, ngroups, dstate, dtype=dtype, device=device) + C = 0.1 * torch.randn(batch, seq_len, ngroups, dstate, dtype=dtype, device=device) + + y_ssd, final_ssd = ssd_minimal_discrete(x * dt.unsqueeze(-1), A * dt, B, C, chunk_size) + + if not vary_A: + # NOTE: fused kernel does not support varying A with time + y_fuse, final_fuse = mamba_chunk_scan_combined(x, dt, A, B, C, chunk_size, D=None, return_final_states=True) + assert y_ssd.allclose(y_fuse, 0, atol), f'y diff: {torch.abs(y_ssd - y_fuse).max()}' + # fused kernel upcasts state to float32 + # https://github.com/state-spaces/mamba/blob/v2.2.2/mamba_ssm/ops/triton/ssd_combined.py#L650 + final_fuse = final_fuse.to(dtype) + assert final_ssd.allclose(final_fuse, 0, atol), f'final diff: {torch.abs(final_ssd - final_fuse).max()}' + + # mapping inputs Mamba2 -> FLA + # FLA Now use head_first = False, therefore there is no need to transpose inputs + q = C + k = B + v = x + g = (A * dt) + + # mapping outputs Mamba2 -> FLA + y_rearrange = y_ssd + final_rearrange = final_ssd.transpose(2, 3) + + # comparing output results between FLA kernel and Mamba2 kernel + # final_gla_fuse :[N, H, K, V] + outputs_gla_fuse, final_gla_fuse = chunk_simple_gla(q, k, v, g, scale=1.0, output_final_state=True) + assert y_rearrange.allclose(outputs_gla_fuse, 0, atol), f'y diff: {torch.abs(y_rearrange - outputs_gla_fuse).max()}' + final_gla_fuse = final_gla_fuse.to(dtype) # states hard-coded to float32 in FLA kernel + assert final_rearrange.allclose(final_gla_fuse, 0, atol), f'final diff: {torch.abs(final_ssd - final_gla_fuse).max()}' diff --git a/code/flash-linear-attention/tests/ops/test_solve_tril.py b/code/flash-linear-attention/tests/ops/test_solve_tril.py new file mode 100644 index 0000000000000000000000000000000000000000..1c0480b9965e0e852315f8c13be47a4bbb038f6f --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_solve_tril.py @@ -0,0 +1,91 @@ + +import os + +import pytest +import torch +import torch.nn.functional as F + +from fla.ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd +from fla.ops.utils.solve_tril import solve_tril +from fla.utils import assert_close, device, device_platform + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'chunk_size'), + [ + pytest.param(*test, id="B{}-T{}-H{}-chunk_size{}".format(*test)) + for test in [ + (1, 63, 1, 16), + (2, 500, 4, 32), + (2, 1000, 5, 64), + (3, 1024, 6, 64), + (4, 2048, 8, 64), + ] + ], +) +@pytest.mark.skipif( + device_platform == 'intel', + reason='Intel Pytorch Failure', +) +def test_solve_tril(B, T, H, chunk_size): + # do not randomly intiialize A otherwise the inverse is not stable + k = F.normalize(torch.randn((B, H, T, 64), dtype=torch.float32, device=device), dim=-1) + # Pad the second-to-last dimension (T) to be a multiple of chunk_size + padding_size = (chunk_size - T % chunk_size) % chunk_size + k_padded = F.pad(k, (0, 0, 0, padding_size, 0, 0, 0, 0)) + k_padded = k_padded.reshape(B, H, -1, chunk_size, 64) + A = (k_padded @ k_padded.transpose(-1, -2)).tril(-1) + + ref = torch.inverse(A + torch.eye(A.shape[-1], device=A.device)[None, None, None, ...]) + ref = ref.reshape(B, H, -1, chunk_size)[:, :, :T, :] + + tri = solve_tril(A.reshape(B, H, -1, chunk_size)[:, :, :T, :].transpose(1, 2)).transpose(1, 2) + + assert_close('solve_tril', ref, tri, 0.0001) + + +@pytest.mark.parametrize( + ('H', 'D', 'chunk_size', 'cu_seqlens'), + [ + pytest.param(*test, id="H{}-D{}-chunk_size{}-cu_seqlens{}".format(*test)) + for test in [ + (4, 64, 16, [0, 15]), + (4, 64, 32, [0, 256, 500, 1000]), + (4, 100, 64, [0, 15, 100, 300, 1200, 2000]), + (4, 64, 16, [0, 1, 100, 300, 1200, 2048]), + (4, 128, 32, [0, 200, 512, 1200, 2048]), + ] + ], +) +@pytest.mark.skipif( + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set', +) +@pytest.mark.skipif( + device_platform == 'intel', + reason='Intel Pytorch Failure', +) +def test_solve_tril_varlen( + H: int, + D: int, + chunk_size: int, + cu_seqlens: list[int], +): + T = cu_seqlens[-1] + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + # Construct the input. otherwise inverse's condition number might be too large to measure the error + k = F.normalize(torch.randn((1, T, H, D), dtype=torch.bfloat16, device=device), dim=-1) + beta = torch.randn((1, T, H), dtype=torch.bfloat16, device=device).sigmoid() + A = chunk_scaled_dot_kkt_fwd(k=k, beta=beta, cu_seqlens=cu_seqlens, chunk_size=chunk_size) + + ref = torch.zeros_like(A) + for i in range(len(cu_seqlens) - 1): + for j in range(cu_seqlens[i], cu_seqlens[i+1], chunk_size): + actual_size = min(chunk_size, cu_seqlens[i+1] - j) + ref[:, j:j+actual_size, :, :actual_size] = torch.inverse( + A[:, j:j+actual_size, :, :actual_size].transpose(1, 2) + + torch.eye(actual_size, device=A.device, dtype=A.dtype)[None, None, ...], + ).transpose(1, 2) + + tri = solve_tril(A, cu_seqlens=cu_seqlens) + assert_close('solve_tril_varlen', ref, tri, 0.0001) diff --git a/code/flash-linear-attention/tests/ops/test_titans.py b/code/flash-linear-attention/tests/ops/test_titans.py new file mode 100644 index 0000000000000000000000000000000000000000..1882eed4a9fb218419e9bedda0f5f73de5d7211b --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_titans.py @@ -0,0 +1,122 @@ + +import pytest +import torch +import torch.nn.functional as F + +from fla.ops.titans.naive import chunk_titans_linear_ref +from fla.utils import assert_close, device + + +def initialize_chunked_param(B, H, T, BT, dtype=torch.float32): + # Calculate number of complete chunks and remaining elements + num_complete_chunks = T // BT + remainder = T % BT + + # Initialize for complete chunks + if num_complete_chunks > 0: + theta_chunks = torch.rand(B, H, num_complete_chunks, 1, dtype=dtype) + theta_main = theta_chunks.repeat_interleave( + BT, dim=2, + ) # Shape: (B, H, num_complete_chunks*BT, 1) + else: + theta_main = torch.empty(B, H, 0, 1, dtype=dtype) + + # Handle remaining elements if any + if remainder > 0: + theta_remainder = torch.rand(B, H, 1, 1, dtype=dtype) + theta_remainder = theta_remainder.repeat_interleave( + remainder, dim=2, + ) # Shape: (B, H, remainder, 1) + + # Concatenate main chunks with remainder + theta = torch.cat([theta_main, theta_remainder], dim=2) + else: + theta = theta_main + + return theta + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (1, 63, 1, 64, torch.float16), + (2, 100, 4, 60, torch.float16), + (2, 1024, 3, 128, torch.float16), + (3, 2000, 4, 128, torch.float16), + (4, 2048, 8, 64, torch.float16), + ] + ], +) +@pytest.mark.skipif( + True, reason='FIXME', +) +def test_naive_chunk( + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): + BT = 64 + # set seed + torch.manual_seed(1) + # we don't use such initialization in the original code + # theta = initialize_chunked_param(B, H, T, BT, dtype) + # alpha = initialize_chunked_param(B, H, T, BT, dtype) + # eta = initialize_chunked_param(B, H, T, BT, dtype) + theta = torch.rand(B, H, T, 1, dtype=dtype) + alpha = torch.rand(B, H, T, 1, dtype=dtype) + eta = torch.rand(B, H, T, 1, dtype=dtype) + + # titans normalize queries and keys using ℓ2-normalization + q = F.normalize(torch.randn(B, H, T, D, dtype=torch.float32), p=2, dim=-1).to(dtype) + k = F.normalize(torch.randn(B, H, T, D, dtype=torch.float32), p=2, dim=-1).to(dtype) + v = torch.randn(B, H, T, D, dtype=dtype) + w = torch.randn(H, D, dtype=dtype) + b = torch.randn(H, D, dtype=dtype) + h0 = torch.randn(B, H, D, D, dtype=torch.float32) + q = q.permute(0, 2, 1, 3) + k = k.permute(0, 2, 1, 3) + v = v.permute(0, 2, 1, 3) + theta = theta.permute(0, 2, 1, 3) + alpha = alpha.permute(0, 2, 1, 3) + eta = eta.permute(0, 2, 1, 3) + q, k, v, w, b, theta, alpha, eta = map( + lambda x: x.to(device).requires_grad_(False), (q, k, v, w, b, theta, alpha, eta), + ) + # in titans paper, h0 is not learnable + h0 = h0.to(device) + + ref_naive, ref_ht_naive = chunk_titans_linear_ref( + q.clone(), + k.clone(), + v.clone(), + w.clone(), + b.clone(), + theta.clone(), + alpha.clone(), + eta.clone(), + output_final_state=True, + chunk_size=BT, + initial_state=h0.clone(), + use_chunk=False, + ) + ref, ref_ht = chunk_titans_linear_ref( + q.clone(), + k.clone(), + v.clone(), + w.clone(), + b.clone(), + theta.clone(), + alpha.clone(), + eta.clone(), + output_final_state=True, + chunk_size=BT, + initial_state=h0.clone(), + use_chunk=True, + ) + + assert_close(" o", ref, ref_naive, 0.006) + assert_close("ht", ref_ht, ref_ht_naive, 0.005) diff --git a/code/flash-linear-attention/tests/ops/test_ttt.py b/code/flash-linear-attention/tests/ops/test_ttt.py new file mode 100644 index 0000000000000000000000000000000000000000..142446c087b541c1269aac35b0aa73363d3638c6 --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_ttt.py @@ -0,0 +1,268 @@ + +import os + +import pytest +import torch +import torch.nn.functional as F + +from fla.ops.ttt import chunk_ttt_linear, fused_chunk_ttt_linear +from fla.ops.ttt.naive import chunk_ttt_linear_ref +from fla.utils import assert_close, check_shared_mem, device + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'scale', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-{}".format(*test)) + for test in [ + (1, 63, 1, 64, 1, torch.float16), + (2, 100, 4, 60, 0.1, torch.float16), + (2, 1024, 3, 128, 0.1, torch.float16), + (2, 1024, 4, 128, 1, torch.float16), + (3, 2000, 4, 128, 0.1, torch.float16), + (4, 2048, 8, 64, 0.1, torch.float16), + ] + ], +) +def test_chunk( + B: int, + T: int, + H: int, + D: int, + scale: float, + dtype: torch.dtype, +): + if D > 64 and check_shared_mem('hopper') is False: + pytest.skip(reason="Current CI do not support this config") + if T > 1000: + pytest.skip(reason="Current CI do not support this config") + eta_base = 5e-3 + q = torch.randn(B, T, H, D, dtype=dtype) + k = F.normalize(torch.randn(B, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype) + v = torch.randn(B, T, H, D, dtype=dtype) + w = torch.randn(H, D, dtype=dtype) + b = torch.randn(H, D, dtype=dtype) + eta = torch.randn(B, T, H, 1, dtype=dtype) * eta_base + h0 = torch.randn(B, H, D, D, dtype=torch.float32) + hb0 = torch.randn(B, H, 1, D, dtype=torch.float32) + + q, k, v, w, b, eta, h0, hb0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, w, b, eta, h0, hb0)) + do = torch.rand_like(v) + dht = torch.rand_like(h0) + dhbt = torch.rand_like(hb0) + + tri, tri_ht, tri_hbt = chunk_ttt_linear( + q.clone(), + k.clone(), + v.clone(), + w.clone(), + b.clone(), + eta.clone(), + scale=scale, + output_final_state=True, + initial_state=h0.clone(), + initial_state_bias=hb0.clone(), + ) + ((tri * do).sum() + (tri_ht * dht).sum() + (tri_hbt * dhbt).sum()).backward(retain_graph=True) + tri_dq, tri_dk, tri_dv, tri_dw, tri_db, tri_deta, \ + tri_dh0, tri_dhb0 = q.grad, k.grad, v.grad, w.grad, b.grad, eta.grad, h0.grad, hb0.grad + q.grad = k.grad = v.grad = w.grad = b.grad = eta.grad = h0.grad = hb0.grad = None + + ref, ref_ht, ref_hbt = chunk_ttt_linear_ref( + q.clone(), + k.clone(), + v.clone(), + w.clone(), + b.clone(), + eta.clone(), + scale=scale, + output_final_state=True, + initial_state=h0.clone(), + initial_state_bias=hb0.clone(), + ) + ((ref * do).sum() + (ref_ht * dht).sum() + (ref_hbt * dhbt).sum()).backward(retain_graph=True) + ref_dq, ref_dk, ref_dv, ref_dw, ref_db, ref_deta, \ + ref_dh0, ref_dhb0 = q.grad, k.grad, v.grad, w.grad, b.grad, eta.grad, h0.grad, hb0.grad + + assert_close(" o", ref, tri, 0.005) + assert_close(" ht", ref_ht, tri_ht, 0.005) + assert_close(" hbt", ref_hbt, tri_hbt, 0.005) + assert_close(" dq", ref_dq, tri_dq, 0.005) + assert_close(" dk", ref_dk, tri_dk, 0.010) + assert_close(" dv", ref_dv, tri_dv, 0.007) + assert_close(" dw", ref_dw, tri_dw, 0.006) + assert_close(" db", ref_db, tri_db, 0.006) + assert_close(" de", ref_deta, tri_deta, 0.030) # because the last element of the chunk + assert_close(" de0", ref_deta[:, :14, :, :], tri_deta[:, :14, :, :], 0.010) + assert_close(" dh0", ref_dh0, tri_dh0, 0.007) + assert_close("dhb0", ref_dhb0, tri_dhb0, 0.005) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'scale', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-scale{}-{}".format(*test)) + for test in [ + (1, 63, 1, 64, 1, torch.float16), + (2, 100, 4, 60, 0.1, torch.float16), + (2, 1024, 3, 128, 0.1, torch.float16), + (2, 1024, 4, 128, 1, torch.float16), + (3, 2000, 4, 128, 0.1, torch.float16), + (4, 2048, 8, 64, 0.1, torch.float16), + ] + ], +) +def test_fused_chunk( + B: int, + T: int, + H: int, + D: int, + scale: float, + dtype: torch.dtype, +): + if D > 64 and check_shared_mem('hopper') is False: + pytest.skip(reason="Current CI do not support this config") + if T > 1000: + pytest.skip(reason="Current CI do not support this config") + eta_base = 5e-3 + q = torch.randn(B, T, H, D, dtype=dtype) + k = F.normalize(torch.randn(B, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype) + v = torch.randn(B, T, H, D, dtype=dtype) + w = torch.randn(H, D, dtype=dtype) + b = torch.randn(H, D, dtype=dtype) + eta = torch.randn(B, T, H, 1, dtype=dtype) * eta_base + h0 = torch.randn(B, H, D, D, dtype=torch.float32) + hb0 = torch.randn(B, H, 1, D, dtype=torch.float32) + + q, k, v, w, b, eta, h0, hb0 = map(lambda x: x.to(device).requires_grad_(True), (q, k, v, w, b, eta, h0, hb0)) + do = torch.rand_like(v) + dht = torch.rand_like(h0) + dhbt = torch.rand_like(hb0) + + tri, tri_ht, tri_hbt = fused_chunk_ttt_linear( + q.clone(), + k.clone(), + v.clone(), + w.clone(), + b.clone(), + eta.clone(), + scale=scale, + output_final_state=True, + initial_state=h0.clone(), + initial_state_bias=hb0.clone(), + ) + ((tri * do).sum() + (tri_ht * dht).sum() + (tri_hbt * dhbt).sum()).backward(retain_graph=True) + tri_dq, tri_dk, tri_dv, tri_dw, tri_db, tri_deta, \ + tri_dh0, tri_dhb0 = q.grad, k.grad, v.grad, w.grad, b.grad, eta.grad, h0.grad, hb0.grad + q.grad = k.grad = v.grad = w.grad = b.grad = eta.grad = h0.grad = hb0.grad = None + + ref, ref_ht, ref_hbt = chunk_ttt_linear_ref( + q.clone(), + k.clone(), + v.clone(), + w.clone(), + b.clone(), + eta.clone(), + scale=scale, + output_final_state=True, + initial_state=h0.clone(), + initial_state_bias=hb0.clone(), + ) + ((ref * do).sum() + (ref_ht * dht).sum() + (ref_hbt * dhbt).sum()).backward(retain_graph=True) + ref_dq, ref_dk, ref_dv, ref_dw, ref_db, ref_deta, \ + ref_dh0, ref_dhb0 = q.grad, k.grad, v.grad, w.grad, b.grad, eta.grad, h0.grad, hb0.grad + + assert_close(" o", ref, tri, 0.005) + assert_close(" ht", ref_ht, tri_ht, 0.005) + assert_close(" hbt", ref_hbt, tri_hbt, 0.005) + assert_close(" dq", ref_dq, tri_dq, 0.005) + assert_close(" dk", ref_dk, tri_dk, 0.010) + assert_close(" dv", ref_dv, tri_dv, 0.007) + assert_close(" dw", ref_dw, tri_dw, 0.005) + assert_close(" db", ref_db, tri_db, 0.005) + assert_close(" de", ref_deta, tri_deta, 0.03) # because the last element of the chunk + assert_close(" de0", ref_deta[:, :14, :, :], tri_deta[:, :14, :, :], 0.008) + assert_close(" dh0", ref_dh0, tri_dh0, 0.006) + assert_close("dhb0", ref_dhb0, tri_dhb0, 0.005) + + +@pytest.mark.parametrize( + ('H', 'D', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-D{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (2, 64, [0, 15], torch.float16), + (3, 60, [0, 111, 500], torch.float16), + (3, 64, [0, 256, 500, 900, 1000], torch.float16), + (4, 100, [0, 15, 100, 300, 1200, 1599, 1800, 2000], torch.float16), + ] + ], +) +@pytest.mark.skipif( + os.getenv("SKIP_TEST_CHUNK_VARLEN") == "1", + reason="Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set", +) +def test_chunk_varlen( + H: int, + D: int, + cu_seqlens: list[int], + dtype: torch.dtype, +): + if D > 64 and check_shared_mem('hopper') is False: + pytest.skip(reason="Current CI do not support this config") + torch.manual_seed(42) + os.environ['TRITON_F32_DEFAULT'] = 'ieee' + T = cu_seqlens[-1] + N = len(cu_seqlens) - 1 + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + + eta_base = 5e-3 + # seq-first required for inputs with variable lengths + q = torch.randn((1, T, H, D), dtype=dtype) + k = F.normalize(torch.randn(1, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype) + v = torch.randn((1, T, H, D), dtype=dtype) + eta = torch.randn(1, T, H, 1, dtype=dtype) * eta_base + w = torch.randn(H, D, dtype=dtype) + b = torch.randn(H, D, dtype=dtype) + h0 = torch.randn((N, H, D, D), dtype=torch.float32) + hb0 = torch.randn((N, H, 1, D), dtype=torch.float32) + q, k, v, w, b, eta, h0, hb0 = map(lambda x: x.to(device).requires_grad_(), (q, k, v, w, b, eta, h0, hb0)) + + tri, tri_ht, tri_hbt = chunk_ttt_linear( + q.clone(), + k.clone(), + v.clone(), + w.clone(), + b.clone(), + eta.clone(), + output_final_state=True, + initial_state=h0.clone(), + initial_state_bias=hb0.clone(), + cu_seqlens=cu_seqlens, + ) + + ref = [] + ref_ht = [] + ref_hbt = [] + for i in range(N): + ref_i, ref_ht_i, ref_hbt_i = chunk_ttt_linear_ref( + q=q[:, cu_seqlens[i]:cu_seqlens[i+1]], + k=k[:, cu_seqlens[i]:cu_seqlens[i+1]], + v=v[:, cu_seqlens[i]:cu_seqlens[i+1]], + w=w, + b=b, + eta=eta[:, cu_seqlens[i]:cu_seqlens[i+1]], + initial_state=h0[i], + initial_state_bias=hb0[i], + output_final_state=True, + ) + ref.append(ref_i) + ref_ht.append(ref_ht_i) + ref_hbt.append(ref_hbt_i) + ref = torch.cat(ref, 1) + ref_ht = torch.cat(ref_ht, 0) + ref_hbt = torch.cat(ref_hbt, 0) + + assert_close(" o", ref, tri, 0.005) + assert_close(" ht", ref_ht, tri_ht, 0.005) + assert_close("hbt", ref_hbt, tri_hbt, 0.005) diff --git a/code/flash-linear-attention/tests/ops/test_utils.py b/code/flash-linear-attention/tests/ops/test_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..cf113dfa716bfed20098338f2f008a572bc692ce --- /dev/null +++ b/code/flash-linear-attention/tests/ops/test_utils.py @@ -0,0 +1,410 @@ + +import os + +import pytest +import torch + +from fla.ops.utils import chunk_global_cumsum, chunk_local_cumsum, mean_pooling +from fla.ops.utils.index import prepare_lens +from fla.ops.utils.pack import pack_sequence, unpack_sequence +from fla.utils import assert_close, device + + +def reversed_cumsum(x, dim=-1): + dtype = x.dtype + x = x.float() + c = x.cumsum(dim) + y = x + c.index_select(dim, x.new_tensor([c.shape[dim]-1], dtype=torch.long)) - c + return y.to(dtype) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (1, 63, 1, 30, torch.float), + (2, 500, 4, 60, torch.float), + (2, 1000, 5, 128, torch.float), + (3, 1024, 6, 500, torch.float), + (4, 2048, 8, 1024, torch.float), + ] + ], +) +def test_global_cumsum( + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): + torch.manual_seed(42) + s = torch.randn(B, T, H, dtype=dtype).to(device) + ref = s.float().cumsum(1).to(dtype) + tri = chunk_global_cumsum(s) + assert_close('global_cumsum', ref, tri, 1e-3) + + s = torch.randn(B, T, H, D, dtype=dtype).to(device) + ref = s.float().cumsum(1).to(dtype) + tri = chunk_global_cumsum(s) + assert_close('global_cumsum', ref, tri, 1e-3) + + +@pytest.mark.parametrize( + ('H', 'D', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-D{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (2, 60, [0, 15], torch.float), + (3, 100, [0, 256, 500, 1000], torch.float), + (4, 256, [0, 15, 100, 300, 1200, 2000], torch.float), + (4, 500, [0, 1, 100, 300, 1200, 2048], torch.float16), + (2, 1024, [0, 200, 512, 1200, 2048], torch.float16), + ] + ], +) +@pytest.mark.skipif( + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set', +) +def test_global_cumsum_varlen( + H: int, + D: int, + cu_seqlens: list[int], + dtype: torch.dtype, +): + torch.manual_seed(42) + T = cu_seqlens[-1] + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + + s = torch.randn(1, T, H, dtype=dtype).to(device) + ref = torch.cat([s[:, start:end].float().cumsum(1) for start, end in zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False)], 1).to(dtype) + tri = chunk_global_cumsum(s, cu_seqlens=cu_seqlens) + assert_close('global_cumsum', ref, tri, 1e-3) + + s = torch.randn(1, T, H, D, dtype=dtype).to(device) + ref = torch.cat([s[:, start:end].float().cumsum(1) for start, end in zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False)], 1).to(dtype) + tri = chunk_global_cumsum(s, cu_seqlens=cu_seqlens) + assert_close('global_cumsum', ref, tri, 1e-3) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (1, 63, 1, 30, torch.float), + (2, 500, 4, 60, torch.float), + (2, 1000, 5, 128, torch.float), + (3, 1024, 6, 500, torch.float), + (4, 2048, 8, 1024, torch.float), + ] + ], +) +def test_global_reversed_cumsum( + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): + torch.manual_seed(42) + s = torch.randn(B, T, H, dtype=dtype).to(device) + ref = reversed_cumsum(s, dim=(1)).to(dtype) + tri = chunk_global_cumsum(s, reverse=True) + assert_close('global_cumsum', ref, tri, 1e-3) + + s = torch.randn(B, T, H, D, dtype=dtype).to(device) + ref = reversed_cumsum(s, dim=(1)).to(dtype) + tri = chunk_global_cumsum(s, reverse=True) + assert_close('global_cumsum', ref, tri, 1e-3) + + +@pytest.mark.parametrize( + ('H', 'D', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-D{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (2, 60, [0, 15], torch.float), + (3, 100, [0, 256, 500, 1000], torch.float), + (4, 256, [0, 15, 100, 300, 1200, 2000], torch.float), + (4, 500, [0, 1, 100, 300, 1200, 2048], torch.float16), + (2, 1024, [0, 200, 512, 1200, 2048], torch.float16), + ] + ], +) +@pytest.mark.skipif( + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set', +) +def test_global_reversed_cumsum_varlen( + H: int, + D: int, + cu_seqlens: list[int], + dtype: torch.dtype, +): + torch.manual_seed(42) + T = cu_seqlens[-1] + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + + s = torch.randn(1, T, H, dtype=dtype).to(device) + ref = torch.cat([reversed_cumsum(s[:, start:end], 1) for start, end in zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False)], 1).to(dtype) + tri = chunk_global_cumsum(s, reverse=True, cu_seqlens=cu_seqlens) + assert_close('global_reversed_cumsum', ref, tri, 1e-3) + + s = torch.randn(1, T, H, D, dtype=dtype).to(device) + ref = torch.cat([reversed_cumsum(s[:, start:end], 1) for start, end in zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False)], 1).to(dtype) + tri = chunk_global_cumsum(s, reverse=True, cu_seqlens=cu_seqlens) + assert_close('global_reversed_cumsum', ref, tri, 1e-3) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'C', 'D', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-C{}-D{}-{}".format(*test)) + for test in [ + (1, 63, 1, 16, 30, torch.float), + (2, 500, 4, 32, 60, torch.float), + (2, 1000, 5, 64, 128, torch.float), + (3, 1024, 6, 64, 500, torch.float), + (4, 2048, 8, 128, 1024, torch.float), + ] + ], +) +def test_local_cumsum( + B: int, + T: int, + H: int, + C: int, + D: int, + dtype: torch.dtype, +): + torch.manual_seed(42) + s = torch.randn(B, T, H, dtype=dtype).to(device) + ref = torch.cat([s[:, i:i+C, :].float().cumsum(1) for i in range(0, T, C)], 1) + tri = chunk_local_cumsum(s, chunk_size=C) + assert_close('local_cumsum', ref, tri, 1e-3) + + s = torch.randn(B, T, H, D, dtype=dtype).to(device) + ref = torch.cat([s[:, i:i+C, :].float().cumsum(1) for i in range(0, T, C)], 1) + tri = chunk_local_cumsum(s, chunk_size=C) + assert_close('local_cumsum', ref, tri, 1e-3) + + +@pytest.mark.parametrize( + ('H', 'C', 'D', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-C{}-D{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (2, 32, 60, [0, 15], torch.float), + (3, 64, 100, [0, 256, 500, 1000], torch.float), + (4, 64, 256, [0, 15, 100, 300, 1200, 2000], torch.float), + (4, 128, 500, [0, 1, 100, 300, 1200, 2048], torch.float16), + (2, 128, 1024, [0, 200, 512, 1200, 2048], torch.float16), + ] + ], +) +@pytest.mark.skipif( + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set', +) +def test_local_cumsum_varlen( + H: int, + C: int, + D: int, + cu_seqlens: list[int], + dtype: torch.dtype, +): + torch.manual_seed(42) + T = cu_seqlens[-1] + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + + s = torch.randn(1, T, H, dtype=dtype).to(device) + ref = torch.cat([ + torch.cat([s[:, i:min(end, i+C), :].float().cumsum(1) for i in range(start, end, C)], 1) + for start, end in zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False) + ], 1) + tri = chunk_local_cumsum(s, chunk_size=C, cu_seqlens=cu_seqlens) + assert_close('local_cumsum', ref, tri, 1e-3) + + s = torch.randn(1, T, H, D, dtype=dtype).to(device) + ref = torch.cat([ + torch.cat([s[:, i:min(end, i+C), :].float().cumsum(1) for i in range(start, end, C)], 1) + for start, end in zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False) + ], 1) + tri = chunk_local_cumsum(s, chunk_size=C, cu_seqlens=cu_seqlens) + assert_close('local_cumsum', ref, tri, 1e-3) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'C', 'D', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-C{}-D{}-{}".format(*test)) + for test in [ + (1, 63, 1, 16, 30, torch.float), + (2, 500, 4, 32, 60, torch.float), + (2, 1000, 5, 64, 128, torch.float), + (3, 1024, 6, 64, 500, torch.float), + (4, 2048, 8, 128, 1024, torch.float), + ] + ], +) +def test_mean_pooling( + B: int, + T: int, + H: int, + C: int, + D: int, + dtype: torch.dtype, +): + torch.manual_seed(42) + x = torch.randn(B, T, H, D, dtype=dtype).to(device) + x.requires_grad = True + ref = torch.cat([x[:, i:i+C, :].float().mean(1, True) for i in range(0, T, C)], 1).to(dtype) + do = torch.randn_like(ref) + ref.backward(do) + ref_dx, x.grad = x.grad.clone(), None + + tri = mean_pooling(x, chunk_size=C) + tri.backward(do) + tri_dx, x.grad = x.grad.clone(), None + + assert_close('mean_pooling', ref, tri, 1e-3) + assert_close('mean_pooling', ref_dx, tri_dx, 1e-3) + + +@pytest.mark.parametrize( + ('H', 'C', 'D', 'cu_seqlens', 'dtype'), + [ + pytest.param(*test, id="H{}-C{}-D{}-cu_seqlens{}-{}".format(*test)) + for test in [ + (2, 32, 60, [0, 15], torch.float), + (3, 64, 100, [0, 256, 500, 1000], torch.float), + (4, 64, 256, [0, 15, 100, 300, 1200, 2000], torch.float), + (4, 128, 500, [0, 1, 100, 300, 1200, 2048], torch.float16), + (2, 128, 1024, [0, 200, 512, 1200, 2048], torch.float16), + ] + ], +) +@pytest.mark.skipif( + os.getenv('SKIP_TEST_CHUNK_VARLEN') == '1', + reason='Skipping test_chunk_varlen because SKIP_TEST_CHUNK_VARLEN is set', +) +def test_mean_pooling_varlen( + H: int, + C: int, + D: int, + cu_seqlens: list[int], + dtype: torch.dtype, +): + torch.manual_seed(42) + T = cu_seqlens[-1] + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + + x = torch.randn(1, T, H, D, dtype=dtype).to(device).requires_grad_(True) + ref = torch.cat([ + torch.cat([x[:, i:min(end, i+C), :].float().mean(1, True) for i in range(start, end, C)], 1) + for start, end in zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False) + ], 1).to(dtype) + do = torch.randn_like(ref) + ref.backward(do) + ref_dx, x.grad = x.grad.clone(), None + + tri = mean_pooling(x, chunk_size=C, cu_seqlens=cu_seqlens) + tri.backward(do) + tri_dx, x.grad = x.grad.clone(), None + + torch.testing.assert_close(ref, tri.to(ref.dtype), rtol=1.6e-2, atol=3e-5) + torch.testing.assert_close(ref_dx, tri_dx.to(ref_dx.dtype), rtol=1.6e-2, atol=3e-5) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'padding_side', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-padding_side{}-{}".format(*test)) + for test in [ + (1, 63, 1, 30, 'left', torch.float), + (2, 500, 4, 60, 'right', torch.float), + (2, 1000, 5, 128, 'left', torch.float), + (3, 1024, 6, 500, 'right', torch.float), + (4, 2048, 8, 1024, 'left', torch.float), + ] + ], +) +def test_pack_sequence( + B: int, + T: int, + H: int, + D: int, + padding_side: str, + dtype: torch.dtype, +): + torch.manual_seed(42) + x = torch.randn(B, T, H, D, dtype=dtype).to(device).requires_grad_(True) + cu_seqlens = torch.cat( + [torch.tensor([0])]+[torch.randint(0, T, (1,)).clamp(min=1) for _ in range(B)], + ).cumsum(-1).to(device) + lens = prepare_lens(cu_seqlens) + + if padding_side == 'left': + ref = torch.cat([x[i, -length:] for i, length in enumerate(lens.tolist())], 0) + else: + ref = torch.cat([x[i, :length] for i, length in enumerate(lens.tolist())], 0) + dy = torch.randn_like(ref) + ref.backward(dy) + ref_dx, x.grad = x.grad.clone(), None + + tri = pack_sequence(x, cu_seqlens, padding_side=padding_side) + tri.backward(dy) + tri_dx, x.grad = x.grad.clone(), None + + assert_close('y', ref, tri, 1e-3) + assert_close('dx', ref_dx, tri_dx, 1e-3) + + +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'padding_side', 'dtype'), + [ + pytest.param(*test, id="B{}-T{}-H{}-D{}-padding_side{}-{}".format(*test)) + for test in [ + (1, 63, 1, 30, 'left', torch.float), + (2, 500, 4, 60, 'right', torch.float), + (2, 1000, 5, 128, 'left', torch.float), + (3, 1024, 6, 500, 'right', torch.float), + (4, 2048, 8, 1024, 'left', torch.float), + ] + ], +) +def test_unpack_sequence( + B: int, + T: int, + H: int, + D: int, + padding_side: str, + dtype: torch.dtype, +): + torch.manual_seed(42) + cu_seqlens = torch.cat( + [torch.tensor([0])]+[torch.randint(0, T, (1,)).clamp(min=1) for _ in range(B)], + ).cumsum(-1).to(device) + lens = prepare_lens(cu_seqlens) + desired_shape = (B, lens.max().item() + torch.randint(0, 10, (1,)).item(), H, D) + + x = torch.randn(cu_seqlens[-1].item(), H, D, dtype=dtype).to(device).requires_grad_(True) + ref = torch.zeros(desired_shape, device=device, dtype=dtype) + dy = torch.randn_like(ref) + for i, (bos, eos) in enumerate(zip(cu_seqlens[:-1].tolist(), cu_seqlens[1:].tolist(), strict=False)): + length = eos - bos + if padding_side == 'left': + ref[i, -length:] = x[bos:eos] + else: + ref[i, :length] = x[bos:eos] + ref.backward(dy) + ref_dx, x.grad = x.grad.clone(), None + + tri = unpack_sequence(x, cu_seqlens, padding_side=padding_side, desired_shape=desired_shape) + tri.backward(dy) + tri_dx, x.grad = x.grad.clone(), None + + assert_close('y', ref, tri, 1e-3) + assert_close('dx', ref_dx, tri_dx, 1e-3) diff --git a/code/flash-linear-attention/utils/convert_from_llama.py b/code/flash-linear-attention/utils/convert_from_llama.py new file mode 100644 index 0000000000000000000000000000000000000000..5b8a2dfb5a5f191820fac94543ea747be49a75fa --- /dev/null +++ b/code/flash-linear-attention/utils/convert_from_llama.py @@ -0,0 +1,189 @@ + +# scripts for converting pretrained hf model weights to fla style +# calling the code to make conversions for mistralai/Mistral-7B-v0.1 would achieve the following results: +# | Tasks |Version|Filter|n-shot| Metric |Value | |Stderr| +# |--------------|------:|------|-----:|----------|-----:|---|-----:| +# |arc_challenge | 1|none | 0|acc |0.5043|± |0.0146| +# | | |none | 0|acc_norm |0.5392|± |0.0146| +# |arc_easy | 1|none | 0|acc |0.8081|± |0.0081| +# | | |none | 0|acc_norm |0.7946|± |0.0083| +# |boolq | 2|none | 0|acc |0.8373|± |0.0065| +# |copa | 1|none | 0|acc |0.9300|± |0.0256| +# |hellaswag | 1|none | 0|acc |0.6127|± |0.0049| +# | | |none | 0|acc_norm |0.8100|± |0.0039| +# |lambada_openai| 1|none | 0|perplexity|3.1810|± |0.0583| +# | | |none | 0|acc |0.7563|± |0.0060| +# |openbookqa | 1|none | 0|acc |0.3260|± |0.0210| +# | | |none | 0|acc_norm |0.4380|± |0.0222| +# |piqa | 1|none | 0|acc |0.8069|± |0.0092| +# | | |none | 0|acc_norm |0.8215|± |0.0089| +# |sciq | 1|none | 0|acc |0.9580|± |0.0063| +# | | |none | 0|acc_norm |0.9390|± |0.0076| +# |winogrande | 1|none | 0|acc |0.7395|± |0.0123| + + +import argparse +import warnings + +import torch +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer + +import fla # noqa + + +def sizeof_fmt(num, suffix='B'): + for unit in ('', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi'): + if abs(num) < 1024.0: + return f'{num:.2f}{unit}{suffix}' + num /= 1024.0 + return f'{num:.2f}Yi{suffix}' + + +def convert( + llama: str, + config: str, + output: str, + precision: str = 'float32', +): + AutoTokenizer.from_pretrained(llama).save_pretrained(output) + llama = AutoModelForCausalLM.from_pretrained(llama, torch_dtype=precision) + print(f"Loading Llama ...\n{llama}") + + config = AutoConfig.from_pretrained(config) + config.torch_dtype = precision + model = AutoModelForCausalLM.from_config(config) + if precision in ['float16', 'fp16']: + model = model.to(torch.float16) + elif precision in ['bfloat16', 'bf16']: + model = model.to(torch.bfloat16) + num_parameters = model.num_parameters() + print(f"Initializing the model from the config:\n{config}\n{model}") + print(f"Number of parameters in total: {num_parameters} ({sizeof_fmt(num_parameters)})") + + print("Copying the weights from Llama to the model ...") + vocab_size = llama.model.embed_tokens.weight.shape[0] + if model.model.embeddings.weight.shape[0] != vocab_size: + warnings.warn(f"Llama and the model have different embedding sizes " + f"({vocab_size} vs {model.model.embeddings.weight.shape[0]}), " + f"the model embeddings will be extended with randomly initialized values or truncated") + vocab_size = min(model.model.embeddings.weight.shape[0], vocab_size) + print("llama.model.embed_tokens -> model.model.embeddings") + model.model.embeddings.weight.data[:vocab_size].copy_(llama.model.embed_tokens.weight[:vocab_size]) + torch.testing.assert_close(model.model.embeddings.weight[:vocab_size], llama.model.embed_tokens.weight[:vocab_size]) + for i in range(config.num_hidden_layers): + if hasattr(model.model.layers[i], 'attn_norm'): + if model.model.layers[i].attn_norm.weight is not None: + print(f"llama.model.layers{i}.input_layernorm.weight -> model.model.layers{i}.attn_norm.weight") + model.model.layers[i].attn_norm.weight.data.copy_(llama.model.layers[i].input_layernorm.weight) + torch.testing.assert_close(model.model.layers[i].attn_norm.weight, + llama.model.layers[i].input_layernorm.weight) + if model.model.layers[i].attn_norm.bias is not None: + print(f"llama.model.layers{i}.input_layernorm.bias -> model.model.layers{i}.attn_norm.bias") + model.model.layers[i].attn_norm.bias.data.copy_(llama.model.layers[i].input_layernorm.bias) + torch.testing.assert_close(model.model.layers[i].attn_norm.bias, + llama.model.layers[i].input_layernorm.bias) + model.model.layers[i].attn_norm.eps = llama.model.layers[i].input_layernorm.variance_epsilon + if hasattr(model.model.layers[i].attn, 'norm'): + if model.model.layers[i].attn.norm.weight is not None: + print(f"llama.model.layers{i}.input_layernorm.weight -> model.model.layers{i}.attn.norm.weight") + model.model.layers[i].attn.norm.weight.data.copy_(llama.model.layers[i].input_layernorm.weight) + torch.testing.assert_close(model.model.layers[i].attn.norm.weight, + llama.model.layers[i].input_layernorm.weight) + if model.model.layers[i].attn.norm.bias is not None: + print(f"llama.model.layers{i}.input_layernorm.bias -> model.model.layers{i}.attn.norm.bias") + model.model.layers[i].attn.norm.bias.data.copy_(llama.model.layers[i].input_layernorm.bias) + torch.testing.assert_close(model.model.layers[i].attn.norm.bias, + llama.model.layers[i].input_layernorm.bias) + model.model.layers[i].attn.norm.eps = llama.model.layers[i].input_layernorm.variance_epsilon + + print(f"llama.model.layers{i}.attn.q_proj.weight -> model.model.layers{i}.attn.q_proj.weight") + model.model.layers[i].attn.q_proj.weight.data.copy_(llama.model.layers[i].self_attn.q_proj.weight) + torch.testing.assert_close(model.model.layers[i].attn.q_proj.weight, llama.model.layers[i].self_attn.q_proj.weight) + if hasattr(llama.model.layers[i].self_attn.q_proj, 'bias') and hasattr(model.model.layers[i].attn.q_proj, 'bias'): + print(f"llama.model.layers{i}.attn.q_proj.bias -> model.model.layers{i}.attn.q_proj.bias") + model.model.layers[i].attn.q_proj.bias.data.copy_(llama.model.layers[i].self_attn.q_proj.bias) + torch.testing.assert_close(model.model.layers[i].attn.q_proj.bias, llama.model.layers[i].self_attn.q_proj.bias) + print(f"llama.model.layers.{i}.attn.k_proj.weight -> model.model.layers.{i}.attn.k_proj.weight") + model.model.layers[i].attn.k_proj.weight.data.copy_(llama.model.layers[i].self_attn.k_proj.weight) + torch.testing.assert_close(model.model.layers[i].attn.k_proj.weight, llama.model.layers[i].self_attn.k_proj.weight) + if hasattr(llama.model.layers[i].self_attn.k_proj, 'bias') and hasattr(model.model.layers[i].attn.k_proj, 'bias'): + print(f"llama.model.layers{i}.attn.k_proj.bias -> model.model.layers{i}.attn.k_proj.bias") + model.model.layers[i].attn.k_proj.bias.data.copy_(llama.model.layers[i].self_attn.k_proj.bias) + torch.testing.assert_close(model.model.layers[i].attn.k_proj.bias, llama.model.layers[i].self_attn.k_proj.bias) + print(f"llama.model.layers.{i}.attn.v_proj.weight -> model.model.layers.{i}.attn.v_proj.weight") + model.model.layers[i].attn.v_proj.weight.data.copy_(llama.model.layers[i].self_attn.v_proj.weight) + torch.testing.assert_close(model.model.layers[i].attn.v_proj.weight, llama.model.layers[i].self_attn.v_proj.weight) + if hasattr(llama.model.layers[i].self_attn.v_proj, 'bias') and hasattr(model.model.layers[i].attn.v_proj, 'bias'): + print(f"llama.model.layers{i}.attn.v_proj.bias -> model.model.layers{i}.attn.v_proj.bias") + model.model.layers[i].attn.v_proj.bias.data.copy_(llama.model.layers[i].self_attn.v_proj.bias) + torch.testing.assert_close(model.model.layers[i].attn.v_proj.bias, llama.model.layers[i].self_attn.v_proj.bias) + + print(f"llama.model.layers.{i}.attn.o_proj.weight -> model.model.layers.{i}.attn.o_proj.weight") + model.model.layers[i].attn.o_proj.weight.data.copy_(llama.model.layers[i].self_attn.o_proj.weight) + torch.testing.assert_close(model.model.layers[i].attn.o_proj.weight, llama.model.layers[i].self_attn.o_proj.weight) + + if hasattr(model.model.layers[i], 'mlp_norm'): + if model.model.layers[i].mlp_norm.weight is not None: + print(f"llama.model.layers{i}.post_attention_layernorm.weight -> model.model.layers{i}.mlp_norm.weight") + model.model.layers[i].mlp_norm.weight.data.copy_(llama.model.layers[i].post_attention_layernorm.weight) + torch.testing.assert_close(model.model.layers[i].mlp_norm.weight, + llama.model.layers[i].post_attention_layernorm.weight) + if model.model.layers[i].mlp_norm.bias is not None: + print(f"llama.model.layers{i}.post_attention_layernorm.bias -> model.model.layers{i}.mlp_norm.bias") + model.model.layers[i].mlp_norm.bias.data.copy_(llama.model.layers[i].post_attention_layernorm.bias) + torch.testing.assert_close(model.model.layers[i].mlp_norm.bias, + llama.model.layers[i].post_attention_layernorm.bias) + model.model.layers[i].mlp_norm.eps = llama.model.layers[i].post_attention_layernorm.variance_epsilon + if hasattr(model.model.layers[i].mlp, 'norm'): + if model.model.layers[i].mlp.norm.weight is not None: + print(f"llama.model.layers{i}.post_attention_layernorm.weight -> model.model.layers{i}.mlp.norm.weight") + model.model.layers[i].mlp.norm.weight.data.copy_(llama.model.layers[i].post_attention_layernorm.weight) + torch.testing.assert_close(model.model.layers[i].mlp.norm.weight, + llama.model.layers[i].post_attention_layernorm.weight) + if model.model.layers[i].mlp.norm.bias is not None: + print(f"llama.model.layers{i}.post_attention_layernorm.bias -> model.model.layers{i}.mlp.norm.bias") + model.model.layers[i].mlp.norm.bias.data.copy_(llama.model.layers[i].post_attention_layernorm.bias) + torch.testing.assert_close(model.model.layers[i].mlp.norm.bias, + llama.model.layers[i].post_attention_layernorm.bias) + model.model.layers[i].mlp.norm.eps = llama.model.layers[i].post_attention_layernorm.variance_epsilon + + print(f"llama.model.layers.{i}.mlp.gate_proj.weight -> model.model.layers.{i}.mlp.gate_proj.weight") + model.model.layers[i].mlp.gate_proj.weight.data.copy_(llama.model.layers[i].mlp.gate_proj.weight) + torch.testing.assert_close(model.model.layers[i].mlp.gate_proj.weight, llama.model.layers[i].mlp.gate_proj.weight) + print(f"llama.model.layers.{i}.mlp.up_proj.weight -> model.model.layers.{i}.mlp.up_proj.weight") + model.model.layers[i].mlp.up_proj.weight.data.copy_(llama.model.layers[i].mlp.up_proj.weight) + torch.testing.assert_close(model.model.layers[i].mlp.up_proj.weight, llama.model.layers[i].mlp.up_proj.weight) + + print(f"llama.model.layers.{i}.mlp.down_proj.weight -> model.model.layers.{i}.mlp.down_proj.weight") + model.model.layers[i].mlp.down_proj.weight.data.copy_(llama.model.layers[i].mlp.down_proj.weight) + torch.testing.assert_close(model.model.layers[i].mlp.down_proj.weight, + llama.model.layers[i].mlp.down_proj.weight) + + if model.model.norm.weight is not None: + print("llama.model.norm.weight -> model.model.norm.weight") + model.model.norm.weight.data.copy_(llama.model.norm.weight) + torch.testing.assert_close(model.model.norm.weight, llama.model.norm.weight) + if model.model.norm.bias is not None: + print("llama.model.norm.bias -> model.model.norm.bias") + model.model.norm.bias.data.copy_(llama.model.norm.bias) + torch.testing.assert_close(model.model.norm.bias, llama.model.norm.bias) + model.model.norm.eps = llama.model.norm.variance_epsilon + + if not model.config.tie_word_embeddings: + print("llama.model.lm_head.weight -> model.lm_head.weight") + model.lm_head.weight.data[:vocab_size].copy_(llama.lm_head.weight[:vocab_size]) + torch.testing.assert_close(model.lm_head.weight[:vocab_size], llama.lm_head.weight[:vocab_size]) + model.config.rope_theta = llama.config.rope_theta + + print(f"Saving converted model to {output} ...\n{model}") + model.save_pretrained(output) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--model", default='mistralai/Mistral-7B-v0.1') + parser.add_argument("--config", default='configs/transformer_7B.json') + parser.add_argument("--output", default='converted/transformer-7B') + parser.add_argument('--precision', type=str, default='float32') + args = parser.parse_args() + convert(args.model, args.config, args.output, precision=args.precision) diff --git a/code/flash-linear-attention/utils/convert_from_rwkv6.py b/code/flash-linear-attention/utils/convert_from_rwkv6.py new file mode 100644 index 0000000000000000000000000000000000000000..6e86b71e80eeb32e376ec552fdd0d5352a57e4dd --- /dev/null +++ b/code/flash-linear-attention/utils/convert_from_rwkv6.py @@ -0,0 +1,213 @@ + +# scripts for converting pretrained hf model weights to fla style +# calling the code to make conversions for RWKV/rwkv-6-world-7b would achieve the following results: +# | Tasks |Version|Filter|n-shot| Metric | Value | |Stderr| +# |--------------|------:|------|-----:|---------------|------:|---|------| +# |arc_challenge | 1|none | 0|acc | 0.4130|± |0.0144| +# | | |none | 0|acc_norm | 0.4403|± |0.0145| +# |arc_easy | 1|none | 0|acc | 0.7382|± |0.0090| +# | | |none | 0|acc_norm | 0.7079|± |0.0093| +# |boolq | 2|none | 0|acc | 0.6823|± |0.0081| +# |copa | 1|none | 0|acc | 0.8700|± |0.0338| +# |hellaswag | 1|none | 0|acc | 0.5508|± |0.0050| +# | | |none | 0|acc_norm | 0.7171|± |0.0045| +# |lambada_openai| 1|none | 0|perplexity | 3.2989|± |0.0634| +# | | |none | 0|acc | 0.7493|± |0.0060| +# |openbookqa | 1|none | 0|acc | 0.3200|± |0.0209| +# | | |none | 0|acc_norm | 0.4440|± |0.0222| +# |piqa | 1|none | 0|acc | 0.7753|± |0.0097| +# | | |none | 0|acc_norm | 0.7894|± |0.0095| +# |sciq | 1|none | 0|acc | 0.9370|± |0.0077| +# | | |none | 0|acc_norm | 0.8860|± |0.0101| +# |winogrande | 1|none | 0|acc | 0.6867|± |0.0130| + +import argparse + +import torch +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer + +import fla # noqa +from fla.utils import device + + +def sizeof_fmt(num, suffix='B'): + for unit in ('', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi'): + if abs(num) < 1024.0: + return f'{num:.2f}{unit}{suffix}' + num /= 1024.0 + return f'{num:.2f}Yi{suffix}' + + +def convert( + rwkv6: str, + config: str, + output: str, +): + torch.manual_seed(1) + AutoTokenizer.from_pretrained(rwkv6, trust_remote_code=True).save_pretrained(output) + rwkv6 = AutoModelForCausalLM.from_pretrained(rwkv6, trust_remote_code=True).to(device) + print(f"Loading rwkv6 ...\n{rwkv6}") + + config = AutoConfig.from_pretrained(config) + model = AutoModelForCausalLM.from_config(config).to(device) + num_parameters = model.num_parameters() + print(f"Initializing the model from the config:\n{config}\n{model}") + print(f"Number of parameters in total: {num_parameters} ({sizeof_fmt(num_parameters)})") + + print("Copying the weights from rwkv6 to the model ...") + print("rwkv6.rwkv.embeddings -> model.model.embeddings") + model.model.embeddings.weight.data.copy_(rwkv6.rwkv.embeddings.weight) + torch.testing.assert_close(model.model.embeddings.weight, rwkv6.rwkv.embeddings.weight) + for i in range(config.num_hidden_layers): + if hasattr(model.model.layers[i], 'pre_norm'): + if model.model.layers[i].pre_norm.weight is not None: + print(f"rwkv6.rwkv.blocks{i}.pre_ln.weight -> model.model.layers{i}.pre_norm.weight") + model.model.layers[i].pre_norm.weight.data.copy_(rwkv6.rwkv.blocks[i].pre_ln.weight) + torch.testing.assert_close(model.model.layers[i].pre_norm.weight, rwkv6.rwkv.blocks[i].pre_ln.weight) + if model.model.layers[i].pre_norm.bias is not None: + print(f"rwkv6.rwkv.blocks{i}.pre_ln.bias -> model.model.layers{i}.pre_norm.bias") + model.model.layers[i].pre_norm.bias.data.copy_(rwkv6.rwkv.blocks[i].pre_ln.bias) + torch.testing.assert_close(model.model.layers[i].pre_norm.bias, rwkv6.rwkv.blocks[i].pre_ln.bias) + model.model.layers[i].pre_norm.eps = rwkv6.rwkv.blocks[i].pre_ln.eps + if model.model.layers[i].attn_norm.weight is not None: + print(f"rwkv6.rwkv.blocks{i}.ln1.weight -> model.model.layers{i}.attn_norm.weight") + model.model.layers[i].attn_norm.weight.data.copy_(rwkv6.rwkv.blocks[i].ln1.weight) + torch.testing.assert_close(model.model.layers[i].attn_norm.weight, rwkv6.rwkv.blocks[i].ln1.weight) + if model.model.layers[i].attn_norm.bias is not None: + print(f"rwkv6.rwkv.blocks{i}.ln1.bias -> model.model.layers{i}.attn_norm.bias") + model.model.layers[i].attn_norm.bias.data.copy_(rwkv6.rwkv.blocks[i].ln1.bias) + torch.testing.assert_close(model.model.layers[i].attn_norm.bias, rwkv6.rwkv.blocks[i].ln1.bias) + model.model.layers[i].attn_norm.eps = rwkv6.rwkv.blocks[i].ln1.eps + + print(f"rwkv6.rwkv.blocks{i}.attention.time_maa_x -> model.model.layers.{i}.attn.x_proj0.mu") + model.model.layers[i].attn.x_proj[0].mu.data.copy_(rwkv6.rwkv.blocks[i].attention.time_maa_x.view(-1)) + torch.testing.assert_close(model.model.layers[i].attn.x_proj[0].mu, + rwkv6.rwkv.blocks[i].attention.time_maa_x.view(-1)) + print(f"rwkv6.rwkv.blocks{i}.attention.time_maa_w1.weight -> model.model.layers{i}.attn.x_proj0.linear.weight") + ww, wk, wv, wr, wg = rwkv6.rwkv.blocks[i].attention.time_maa_w1.view(config.hidden_size, 5, -1).unbind(-2) + w = torch.cat((wr, ww, wk, wv, wg), -1).t() + model.model.layers[i].attn.x_proj[0].linear.weight.data.copy_(w) + torch.testing.assert_close(model.model.layers[i].attn.x_proj[0].linear.weight, w) + + print(f"rwkv6.rwkv.blocks{i}.attention.time_maa_w2.weight -> model.model.layers{i}.attn.x_proj2.weight") + ww, wk, wv, wr, wg = rwkv6.rwkv.blocks[i].attention.time_maa_w2.unbind(0) + w = torch.cat((wr, ww, wk, wv, wg), 0).t() + model.model.layers[i].attn.x_proj[2].weight.data.copy_(w) + torch.testing.assert_close(model.model.layers[i].attn.x_proj[2].weight, w) + + print(f"rwkv6.rwkv.blocks{i}.attention.time_maa_wkvrg -> model.model.layers{i}.attn.x_bias") + bias = torch.stack((rwkv6.rwkv.blocks[i].attention.time_maa_r.view(-1), + rwkv6.rwkv.blocks[i].attention.time_maa_w.view(-1), + rwkv6.rwkv.blocks[i].attention.time_maa_k.view(-1), + rwkv6.rwkv.blocks[i].attention.time_maa_v.view(-1), + rwkv6.rwkv.blocks[i].attention.time_maa_g.view(-1))) + model.model.layers[i].attn.x_bias.data.copy_(bias) + torch.testing.assert_close(model.model.layers[i].attn.x_bias, bias) + + print(f"rwkv6.rwkv.blocks{i}.attention.receptance.weight -> model.model.layers{i}.attn.r_proj.linear.weight") + model.model.layers[i].attn.r_proj.linear.weight.data.copy_(rwkv6.rwkv.blocks[i].attention.receptance.weight) + torch.testing.assert_close(model.model.layers[i].attn.r_proj.linear.weight, + rwkv6.rwkv.blocks[i].attention.receptance.weight) + print(f"rwkv6.rwkv.blocks{i}.attention.time_decay_w1 -> model.model.layers{i}.attn.w_proj.linear.lora0.weight") + model.model.layers[i].attn.w_proj.linear.lora[0].weight.data.copy_(rwkv6.rwkv.blocks[i].attention.time_decay_w1.t()) + torch.testing.assert_close(model.model.layers[i].attn.w_proj.linear.lora[0].weight, + rwkv6.rwkv.blocks[i].attention.time_decay_w1.t()) + print(f"rwkv6.rwkv.blocks{i}.attention.time_decay_w2 -> model.model.layers{i}.attn.w_proj.linear.lora2.weight") + model.model.layers[i].attn.w_proj.linear.lora[2].weight.data.copy_(rwkv6.rwkv.blocks[i].attention.time_decay_w2.t()) + torch.testing.assert_close(model.model.layers[i].attn.w_proj.linear.lora[2].weight, + rwkv6.rwkv.blocks[i].attention.time_decay_w2.t()) + print(f"rwkv6.rwkv.blocks{i}.attention.time_decay -> model.model.layers{i}.attn.w_proj.linear.lora2.bias") + model.model.layers[i].attn.w_proj.linear.lora[2].bias.data.copy_(rwkv6.rwkv.blocks[i].attention.time_decay.view(-1)) + torch.testing.assert_close(model.model.layers[i].attn.w_proj.linear.lora[2].bias, + rwkv6.rwkv.blocks[i].attention.time_decay.view(-1)) + + print(f"rwkv6.rwkv.blocks{i}.attention.key.weight -> model.model.layers.{i}.attn.k_proj.linear.weight") + model.model.layers[i].attn.k_proj.linear.weight.data.copy_(rwkv6.rwkv.blocks[i].attention.key.weight) + torch.testing.assert_close(model.model.layers[i].attn.k_proj.linear.weight, + rwkv6.rwkv.blocks[i].attention.key.weight) + print(f"rwkv6.rwkv.blocks{i}.attention.value.weight -> model.model.layers.{i}.attn.v_proj.linear.weight") + model.model.layers[i].attn.v_proj.linear.weight.data.copy_(rwkv6.rwkv.blocks[i].attention.value.weight) + torch.testing.assert_close(model.model.layers[i].attn.v_proj.linear.weight, + rwkv6.rwkv.blocks[i].attention.value.weight) + print(f"rwkv6.rwkv.blocks{i}.attention.gate.weight -> model.model.layers.{i}.attn.g_proj.linear.weight") + model.model.layers[i].attn.g_proj.linear.weight.data.copy_(rwkv6.rwkv.blocks[i].attention.gate.weight) + torch.testing.assert_close(model.model.layers[i].attn.g_proj.linear.weight, + rwkv6.rwkv.blocks[i].attention.gate.weight) + print(f"rwkv6.rwkv.blocks{i}.attention.time_faaaa -> model.model.layers.{i}.attn.bonus") + bonus = rwkv6.rwkv.blocks[i].attention.time_faaaa.view(config.num_heads, -1) + model.model.layers[i].attn.bonus.data.copy_(bonus) + torch.testing.assert_close(model.model.layers[i].attn.bonus, bonus) + + if model.model.layers[i].attn.g_norm.weight is not None: + print(f"rwkv6.rwkv.blocks{i}.attention.ln_x.weight -> model.model.layers[i].attn.g_norm.weight") + model.model.layers[i].attn.g_norm.weight.data.copy_(rwkv6.rwkv.blocks[i].attention.ln_x.weight) + torch.testing.assert_close(model.model.layers[i].attn.g_norm.weight, rwkv6.rwkv.blocks[i].attention.ln_x.weight) + if model.model.layers[i].attn.g_norm.bias is not None: + print(f"rwkv6.rwkv.blocks{i}.attention.ln_x.bias -> model.model.layers[i].attn.g_norm.bias") + model.model.layers[i].attn.g_norm.bias.data.copy_(rwkv6.rwkv.blocks[i].attention.ln_x.bias) + torch.testing.assert_close(model.model.layers[i].attn.g_norm.bias, rwkv6.rwkv.blocks[i].attention.ln_x.bias) + model.model.layers[i].attn.g_norm.eps = rwkv6.rwkv.blocks[i].attention.ln_x.eps + + print(f"rwkv6.rwkv.blocks{i}.attention.output.weight -> model.model.layers.{i}.attn.o_proj.weight") + model.model.layers[i].attn.o_proj.weight.data.copy_(rwkv6.rwkv.blocks[i].attention.output.weight) + torch.testing.assert_close(model.model.layers[i].attn.o_proj.weight, rwkv6.rwkv.blocks[i].attention.output.weight) + + if model.model.layers[i].ffn_norm.weight is not None: + print(f"rwkv6.rwkv.blocks{i}.ln2.weight -> model.model.layers{i}.ffn_norm.weight") + model.model.layers[i].ffn_norm.weight.data.copy_(rwkv6.rwkv.blocks[i].ln2.weight) + torch.testing.assert_close(model.model.layers[i].ffn_norm.weight, rwkv6.rwkv.blocks[i].ln2.weight) + if model.model.layers[i].ffn_norm.bias is not None: + print(f"rwkv6.rwkv.blocks{i}.ln2.bias -> model.model.layers{i}.ffn_norm.bias") + model.model.layers[i].ffn_norm.bias.data.copy_(rwkv6.rwkv.blocks[i].ln2.bias) + torch.testing.assert_close(model.model.layers[i].ffn_norm.bias, rwkv6.rwkv.blocks[i].ln2.bias) + model.model.layers[i].ffn_norm.eps = rwkv6.rwkv.blocks[i].ln2.eps + + print(f"rwkv6.rwkv.blocks{i}.feed_forward.key.weight -> model.model.layers.{i}.ffn.key.linear.weight") + model.model.layers[i].ffn.key.linear.weight.data.copy_(rwkv6.rwkv.blocks[i].feed_forward.key.weight) + torch.testing.assert_close(model.model.layers[i].ffn.key.linear.weight, + rwkv6.rwkv.blocks[i].feed_forward.key.weight) + print(f"rwkv6.rwkv.blocks{i}.feed_forward.time_maa_k -> model.model.layers.{i}.ffn.key.mu") + model.model.layers[i].ffn.key.mu.data.copy_(rwkv6.rwkv.blocks[i].feed_forward.time_maa_k.view(-1)) + torch.testing.assert_close(model.model.layers[i].ffn.key.mu, + rwkv6.rwkv.blocks[i].feed_forward.time_maa_k.view(-1)) + + print(f"rwkv6.rwkv.blocks{i}.feed_forward.value.weight -> model.model.layers.{i}.ffn.value.weight") + model.model.layers[i].ffn.value.weight.data.copy_(rwkv6.rwkv.blocks[i].feed_forward.value.weight) + torch.testing.assert_close(model.model.layers[i].ffn.value.weight, + rwkv6.rwkv.blocks[i].feed_forward.value.weight) + + print(f"rwkv6.rwkv.blocks{i}.feed_forward.receptance.weight -> model.model.layers.{i}.ffn.receptance.linear.weight") + model.model.layers[i].ffn.receptance.linear.weight.data.copy_(rwkv6.rwkv.blocks[i].feed_forward.receptance.weight) + torch.testing.assert_close(model.model.layers[i].ffn.receptance.linear.weight, + rwkv6.rwkv.blocks[i].feed_forward.receptance.weight) + print(f"rwkv6.rwkv.blocks{i}.feed_forward.time_maa_r -> model.model.layers.{i}.ffn.receptance.mu") + model.model.layers[i].ffn.receptance.mu.data.copy_(rwkv6.rwkv.blocks[i].feed_forward.time_maa_r.view(-1)) + torch.testing.assert_close(model.model.layers[i].ffn.receptance.mu, + rwkv6.rwkv.blocks[i].feed_forward.time_maa_r.view(-1)) + + if model.model.norm.weight is not None: + print("rwkv6.rwkv.ln_out.weight -> model.model.norm.weight") + model.model.norm.weight.data.copy_(rwkv6.rwkv.ln_out.weight) + torch.testing.assert_close(model.model.norm.weight, rwkv6.rwkv.ln_out.weight) + if model.model.norm.bias is not None: + print("rwkv6.rwkv.ln_out.bias -> model.model.norm.bias") + model.model.norm.bias.data.copy_(rwkv6.rwkv.ln_out.bias) + torch.testing.assert_close(model.model.norm.bias, rwkv6.rwkv.ln_out.bias) + model.model.norm.eps = rwkv6.rwkv.ln_out.eps + + if not model.config.tie_word_embeddings: + print("rwkv6.rwkv.head.weight -> model.lm_head.weight") + model.lm_head.weight.data.copy_(rwkv6.head.weight) + torch.testing.assert_close(model.lm_head.weight, rwkv6.head.weight) + + print(f"Saving converted model \n{model}\n to {output} ...") + model.save_pretrained(output) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--model", default='RWKV/rwkv-6-world-7b') + parser.add_argument("--config", default='configs/rwkv6_7B.json') + parser.add_argument("--output", default='converted/rwkv6-7B') + args = parser.parse_args() + convert(args.model, args.config, args.output) diff --git a/code/flash-linear-attention/utils/convert_from_rwkv7.py b/code/flash-linear-attention/utils/convert_from_rwkv7.py new file mode 100644 index 0000000000000000000000000000000000000000..4189215ef826212d2d1442b1e2bce63de7d5093f --- /dev/null +++ b/code/flash-linear-attention/utils/convert_from_rwkv7.py @@ -0,0 +1,150 @@ + +# scripts for converting pretrained hf model weights to fla style + +import argparse +import os +import re + +import torch +from transformers import AutoModelForCausalLM + +import fla # noqa +from fla.models.rwkv7 import RWKV7Config + + +def convert( + rwkv7: str, + output: str, + precision: str = 'float32', +): + weights = torch.load(rwkv7, weights_only=True, map_location='cpu') + config = RWKV7Config() + config.vocab_size = weights['emb.weight'].shape[0] # 50304 + config.hidden_size = weights['blocks.0.ffn.key.weight'].shape[1] # 768 + config.hidden_ratio = weights['blocks.0.ffn.key.weight'].shape[0] / weights['blocks.0.ffn.key.weight'].shape[1] # 4.0 + config.intermediate_size = weights['blocks.0.ffn.key.weight'].shape[0] + config.num_hidden_layers = 0 + while f'blocks.{config.num_hidden_layers}.ffn.key.weight' in weights: + config.num_hidden_layers += 1 + # 12 + config.value_dim = [config.hidden_size] * config.num_hidden_layers + config.decay_low_rank_dim = weights['blocks.0.att.w1'].shape[1] # 64 + config.gate_low_rank_dim = weights['blocks.0.att.g1'].shape[1] # 128 + config.a_low_rank_dim = weights['blocks.0.att.a1'].shape[1] # 64 + try: + config.v_low_rank_dim = weights['blocks.1.att.v1'].shape[1] # 32 + except KeyError: + config.v_low_rank_dim = 32 + + if precision in ['bf16', 'bfloat16']: + precision = 'bfloat16' + dtype = torch.bfloat16 + if precision in ['fp16', 'float16']: + precision = 'float16' + dtype = torch.float16 + if precision in ['fp64', 'double', 'float64']: + precision = 'float64' + dtype = torch.float64 + + config.torch_dtype = precision + print(f"Creating model with config:\n{config}") + model = AutoModelForCausalLM.from_config(config).to(dtype=dtype) + + print(model) + model_dict = model.state_dict() + model_names = [n for n in model_dict] + + # these parameters may be present in pth file but are never used: + unused_names = ['blocks.0.att.v0', 'blocks.0.att.v1', 'blocks.0.att.v2'] + # these parameters may or may not be present in pth file: + possible_absent_weights = [ + 'model.layers.0.pre_norm.weight', 'model.layers.0.pre_norm.bias', + ] + # other parameters may raise a KeyError + + def translate_into_fla(name): + transposed = False + emb_head = { + 'emb.weight': 'model.embeddings.weight', + 'ln_out.weight': 'model.norm.weight', + 'ln_out.bias': 'model.norm.bias', + 'head.weight': 'lm_head.weight', + } + proj = { + 'receptance': 'r_proj', + 'key': 'k_proj', + 'value': 'v_proj', + 'ln_x': 'g_norm', + 'output': 'o_proj', + } + if name in unused_names: + return '', False + if name in emb_head: + return emb_head[name], False + name_compo = name.split('.') + assert name_compo[0] == 'blocks' + name_compo[0] = 'model.layers' + assert int(name_compo[1]) in range(config.num_hidden_layers) + name_compo[2] = { + 'att': 'attn', + 'ffn': 'ffn', + 'ln0': 'pre_norm', + 'ln1': 'attn_norm', + 'ln2': 'ffn_norm', + }[name_compo[2]] + if re.match("[wvag][012]", name_compo[3]): + typ, num = name_compo[3] + name_compo[3] = f'{typ}_lora.lora.' + { + '0': '2.bias', + '1': '0.weight', + '2': '2.weight', + }[num] + transposed |= (num in ['1', '2']) + elif name_compo[2] == 'attn' and name_compo[3] in proj: + name_compo[3] = proj[name_compo[3]] + return '.'.join(name_compo), transposed + + for name in weights: + fla_name, transposed = translate_into_fla(name) + print(f'{name:32} -> {fla_name:42}, {transposed}') + if not fla_name: + print('redundant parameters in source weight: ', name, '\n') + continue + weight = weights[name] + # print shape information + shape1 = list(weight.shape) + shape2 = list(model_dict[fla_name].shape) + print(f'{str(shape1):32} {str(shape2)}\n') + + if transposed: + weight.t_() + if shape1 == [1, 1, config.hidden_size]: + weight.squeeze_() + + if "attn.x_" in fla_name: + assert model_dict[fla_name].shape[2:] == weight.shape, \ + f"Shape mismatch for {fla_name}: model_dict={model_dict[fla_name].shape}, weight={weight.shape}" + else: + assert model_dict[fla_name].shape == weight.shape, \ + f"Shape mismatch for {fla_name}: model_dict={model_dict[fla_name].shape}, weight={weight.shape}" + + model_dict[fla_name].data.copy_(weight) + model_names.remove(fla_name) + + print("uninitialized parameters: ", model_names) + for n in model_names: + if n not in possible_absent_weights: + raise KeyError(n) + + os.makedirs(output, exist_ok=True) + + model.save_pretrained(output, max_shard_size="1000GB") + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='Convert RWKV7') + parser.add_argument('--rwkv7', type=str, help='Path to the input model') + parser.add_argument('--output', type=str, help='Directory to save model') + parser.add_argument('--precision', type=str, default='float32') + args = parser.parse_args() + convert(args.rwkv7, args.output, precision=args.precision) diff --git a/code/inference/README.md b/code/inference/README.md new file mode 100644 index 0000000000000000000000000000000000000000..9e447db8badcae3526213b208c39f5b27dfe287a --- /dev/null +++ b/code/inference/README.md @@ -0,0 +1,73 @@ +# Inference Recipes + +Bash-level inference scripts mirroring `train/` — one script per memory row, all calling `inference/unified_inference.py`. + +## Usage + +```bash +export WAN_BASE_MODEL=/path/to/Wan2.1-T2V-1.3B + +# Single memory type +CKPT=./ckpts/context_k1/epoch-0.safetensors \ + bash inference/memory_baselines_basic/run_infer_context_k1.sh + +# With custom prompt and context image +CKPT=./ckpts/context_k1/epoch-0.safetensors \ +PROMPT="A toy bear on a table" \ +CONTEXT_IMAGE=assets/opendomain_revisit/1774363417.png \ + bash inference/memory_baselines_basic/run_infer_context_k1.sh + +# All memory baselines (needs CKPT_DIR with per-row folders) +CKPT_DIR=./ckpts bash inference/memory_baselines_basic/run_infer_all.sh + +# Dynamic SpatialVID row +CKPT=/path/to/retrained_dynamic_spatial_mem/epoch-0.safetensors \ + bash inference/dynamic_spatialvid/run_infer_dyn_spatial_mem.sh +``` + +## Environment Variables + +| Variable | Default | Description | +|---|---|---| +| `CKPT` | (required) | Path to `.safetensors` checkpoint | +| `WAN_BASE_MODEL` | (required) | Wan 2.1 base model directory | +| `PROMPT` | Generic game scene prompt | Text prompt | +| `CONTEXT_IMAGE` | (none) | First-frame context image path | +| `ACTION_PATH` | `env/action_rotation_left_45.json` | Camera trajectory JSON | +| `SEED` | `0` | Random seed | +| `HEIGHT` / `WIDTH` | `352` / `640` | Resolution | +| `NUM_FRAMES` | `81` | Frames per chunk | +| `NUM_INFERENCE_STEPS` | `50` | Denoising steps | +| `SIGMA_SHIFT` | `15.0` (memory baselines) / `5.0` (context learning) | Timestep shift | +| `INFER_OUTPUT_ROOT` | `inference_outputs/` | Output directory | + +## Script Mapping + +### Memory Baselines (`inference/memory_baselines_basic/`) + +| Inference script | `--memory_type` | Training script | +|---|---|---| +| `run_infer_no_memory.sh` | `no_memory` | `run_ablation_no_memory_baseline_two_chunk.sh` | +| `run_infer_framepack_weight.sh` | `framepack_weight` | `run_ablation_framepack_weight_two_chunk.sh` | +| `run_infer_framepack_len_r2.sh` | `framepack_len_r2` | `run_ablation_framepack_len_r2_two_chunk.sh` | +| `run_infer_framepack_len_r4.sh` | `framepack_len_r4` | `run_ablation_framepack_len_r4_two_chunk.sh` | +| `run_infer_framepack_hybrid_r2.sh` | `framepack_hybrid_r2` | `run_ablation_framepack_hybrid_r2_weight_two_chunk.sh` | +| `run_infer_framepack_hybrid_r4.sh` | `framepack_hybrid_r4` | `run_ablation_framepack_hybrid_r4_weight_two_chunk.sh` | +| `run_infer_spatial_mem.sh` | `spatial_mem` | `run_spatial_memory_baseline.sh` | +| `run_infer_spatial_concat_text.sh` | `spatial_concat_text` | `run_ablation_spatial_concat_text_two_chunk.sh` | +| `run_infer_spatial_inject_none.sh` | `spatial_inject_none` | `run_ablation_spatial_inject_none_two_chunk.sh` | +| `run_infer_spatial_cross_attn_readout.sh` | `spatial_cross_attn_readout` | `run_ablation_spatial_cross_attn_readout_two_chunk.sh` | +| `run_infer_videossm_hybrid.sh` | `videossm_hybrid` | `run_videossm_hybrid_baseline.sh` | +| `run_infer_block_wise_ssm.sh` | `block_wise_ssm` | `run_ablation_block_wise_ssm_two_chunk.sh` | + +### Context Learning (`inference/context_learning/`) + +| Inference script | `--memory_type` | Training script | +|---|---|---| +| `run_infer_ctx1.sh` | `context_k1` | `run_pre_qkv_ctx1.sh` | +| `run_infer_ctx5.sh` | `context_k5` | `run_pre_qkv_ctx5.sh` | +| `run_infer_ctx20.sh` | `context_k20` | `run_pre_qkv_ctx20.sh` | + +### Dynamic SpatialVID (`inference/dynamic_spatialvid/`) + +Dynamic wrappers mirror the six dynamic training rows in `train/dynamic_spatialvid/`. They are intended for qualitative replay and demo generation; dynamic evaluation scripts are TODO. diff --git a/code/inference/_shared/common_env_infer.sh b/code/inference/_shared/common_env_infer.sh new file mode 100644 index 0000000000000000000000000000000000000000..e7a8595c57e0c7bc37f593da024936a5efd4ee2e --- /dev/null +++ b/code/inference/_shared/common_env_infer.sh @@ -0,0 +1,59 @@ +#!/bin/bash +# Shared inference environment — mirrors train/_shared/common_env_memory.sh +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[1]:-${BASH_SOURCE[0]}}")" && pwd)" +REPO_ROOT="${REPO_ROOT:-$(cd "${SCRIPT_DIR}/../.." && pwd)}" +export PYTHONPATH="${REPO_ROOT}:${PYTHONPATH:-}" + +# ── Base model ──────────────────────────────────────────────────────── +WAN_BASE_MODEL="${WAN_BASE_MODEL:-}" +if [ -z "${WAN_BASE_MODEL}" ]; then + for _d in \ + "${REPO_ROOT}/checkpoints/Wan2.1-T2V-1.3B" \ + "${REPO_ROOT}/models/Wan2.1-T2V-1.3B" \ + "${REPO_ROOT}/Wan2.1-T2V-1.3B"; do + [ -d "${_d}" ] && { WAN_BASE_MODEL="${_d}"; break; } + done +fi +if [ -z "${WAN_BASE_MODEL}" ]; then + echo "[common_env_infer] ERROR: WAN_BASE_MODEL not set." >&2 + echo "[common_env_infer] HINT: export WAN_BASE_MODEL=/path/to/Wan2.1-T2V-1.3B" >&2 + exit 2 +fi +for _f in diffusion_pytorch_model.safetensors models_t5_umt5-xxl-enc-bf16.pth Wan2.1_VAE.pth; do + [ -f "${WAN_BASE_MODEL}/${_f}" ] || { + echo "[common_env_infer] ERROR: missing ${WAN_BASE_MODEL}/${_f}" >&2; exit 2; } +done +TOKENIZER_PATH="${TOKENIZER_PATH:-}" +if [ -z "${TOKENIZER_PATH}" ] && [ -d "${WAN_BASE_MODEL}/google/umt5-xxl" ]; then + TOKENIZER_PATH="${WAN_BASE_MODEL}/google/umt5-xxl" +fi +export WAN_BASE_MODEL +export TOKENIZER_PATH + +# ── Checkpoint ──────────────────────────────────────────────────────── +# CKPT must be set by the caller or the wrapper script. +# Inference scripts validate this themselves. + +# ── Output ──────────────────────────────────────────────────────────── +INFER_OUTPUT_ROOT="${INFER_OUTPUT_ROOT:-${REPO_ROOT}/inference_outputs}" +mkdir -p "${INFER_OUTPUT_ROOT}" +export INFER_OUTPUT_ROOT + +# ── Defaults ────────────────────────────────────────────────────────── +PROMPT="${PROMPT:-A game scene, the camera moves through the environment}" +CONTEXT_IMAGE="${CONTEXT_IMAGE:-}" +ACTION_PATH="${ACTION_PATH:-${REPO_ROOT}/env/action_rotation_left_45.json}" +SEED="${SEED:-0}" +HEIGHT="${HEIGHT:-352}" +WIDTH="${WIDTH:-640}" +NUM_FRAMES="${NUM_FRAMES:-81}" +NUM_INFERENCE_STEPS="${NUM_INFERENCE_STEPS:-50}" +SIGMA_SHIFT="${SIGMA_SHIFT:-15.0}" +CFG_SCALE="${CFG_SCALE:-5.0}" +FPS="${FPS:-15}" + +echo "[common_env_infer] WAN_BASE_MODEL=${WAN_BASE_MODEL}" +echo "[common_env_infer] INFER_OUTPUT_ROOT=${INFER_OUTPUT_ROOT}" +cd "${REPO_ROOT}" diff --git a/code/inference/context_learning/common_env.sh b/code/inference/context_learning/common_env.sh new file mode 100644 index 0000000000000000000000000000000000000000..30363c2ea921ae79ba6965288e4bea7a10dc3a22 --- /dev/null +++ b/code/inference/context_learning/common_env.sh @@ -0,0 +1,5 @@ +#!/bin/bash +# Common inference wrapper for context_learning. +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SIGMA_SHIFT="${SIGMA_SHIFT:-5.0}" +source "${SCRIPT_DIR}/../_shared/common_env_infer.sh" diff --git a/code/inference/context_learning/run_infer_ctx1.sh b/code/inference/context_learning/run_infer_ctx1.sh new file mode 100644 index 0000000000000000000000000000000000000000..bb4448ff67f386719c4aa1dc1e8fe5a5c12b57f6 --- /dev/null +++ b/code/inference/context_learning/run_infer_ctx1.sh @@ -0,0 +1,28 @@ +#!/bin/bash +# Inference: Context K=1 +# Corresponds to: train/context_learning/run_pre_qkv_ctx1.sh +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" + +: "${CKPT:?ERROR: set CKPT to your context_k1 checkpoint path}" +OUTPUT="${INFER_OUTPUT_ROOT}/context_k1_$(date +%Y%m%d_%H%M%S).mp4" + +EXTRA_ARGS=() +if [ -n "${CONTEXT_IMAGE}" ]; then + EXTRA_ARGS+=(--context_image "${CONTEXT_IMAGE}") +fi +if [ -n "${ACTION_PATH}" ]; then + EXTRA_ARGS+=(--action_path "${ACTION_PATH}") +fi + +python inference/unified_inference.py \ + --ckpt "${CKPT}" \ + --memory_type context_k1 \ + --base_model "${WAN_BASE_MODEL}" \ + --prompt "${PROMPT}" \ + "${EXTRA_ARGS[@]}" \ + --height "${HEIGHT}" --width "${WIDTH}" --num_frames "${NUM_FRAMES}" \ + --seed "${SEED}" --num_inference_steps "${NUM_INFERENCE_STEPS}" \ + --sigma_shift "${SIGMA_SHIFT}" --cfg_scale "${CFG_SCALE}" --fps "${FPS}" \ + --output_path "${OUTPUT}" diff --git a/code/inference/context_learning/run_infer_ctx20.sh b/code/inference/context_learning/run_infer_ctx20.sh new file mode 100644 index 0000000000000000000000000000000000000000..afe41f75e054ce03aa8c6b929541de2a8c3e80f2 --- /dev/null +++ b/code/inference/context_learning/run_infer_ctx20.sh @@ -0,0 +1,28 @@ +#!/bin/bash +# Inference: Context K=20 +# Corresponds to: train/context_learning/run_pre_qkv_ctx20.sh +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" + +: "${CKPT:?ERROR: set CKPT to your context_k20 checkpoint path}" +OUTPUT="${INFER_OUTPUT_ROOT}/context_k20_$(date +%Y%m%d_%H%M%S).mp4" + +EXTRA_ARGS=() +if [ -n "${CONTEXT_IMAGE}" ]; then + EXTRA_ARGS+=(--context_image "${CONTEXT_IMAGE}") +fi +if [ -n "${ACTION_PATH}" ]; then + EXTRA_ARGS+=(--action_path "${ACTION_PATH}") +fi + +python inference/unified_inference.py \ + --ckpt "${CKPT}" \ + --memory_type context_k20 \ + --base_model "${WAN_BASE_MODEL}" \ + --prompt "${PROMPT}" \ + "${EXTRA_ARGS[@]}" \ + --height "${HEIGHT}" --width "${WIDTH}" --num_frames "${NUM_FRAMES}" \ + --seed "${SEED}" --num_inference_steps "${NUM_INFERENCE_STEPS}" \ + --sigma_shift "${SIGMA_SHIFT}" --cfg_scale "${CFG_SCALE}" --fps "${FPS}" \ + --output_path "${OUTPUT}" diff --git a/code/inference/context_learning/run_infer_ctx5.sh b/code/inference/context_learning/run_infer_ctx5.sh new file mode 100644 index 0000000000000000000000000000000000000000..c9e3aa66af082e4b663e20edcc20cf6bfbf42a64 --- /dev/null +++ b/code/inference/context_learning/run_infer_ctx5.sh @@ -0,0 +1,28 @@ +#!/bin/bash +# Inference: Context K=5 +# Corresponds to: train/context_learning/run_pre_qkv_ctx5.sh +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" + +: "${CKPT:?ERROR: set CKPT to your context_k5 checkpoint path}" +OUTPUT="${INFER_OUTPUT_ROOT}/context_k5_$(date +%Y%m%d_%H%M%S).mp4" + +EXTRA_ARGS=() +if [ -n "${CONTEXT_IMAGE}" ]; then + EXTRA_ARGS+=(--context_image "${CONTEXT_IMAGE}") +fi +if [ -n "${ACTION_PATH}" ]; then + EXTRA_ARGS+=(--action_path "${ACTION_PATH}") +fi + +python inference/unified_inference.py \ + --ckpt "${CKPT}" \ + --memory_type context_k5 \ + --base_model "${WAN_BASE_MODEL}" \ + --prompt "${PROMPT}" \ + "${EXTRA_ARGS[@]}" \ + --height "${HEIGHT}" --width "${WIDTH}" --num_frames "${NUM_FRAMES}" \ + --seed "${SEED}" --num_inference_steps "${NUM_INFERENCE_STEPS}" \ + --sigma_shift "${SIGMA_SHIFT}" --cfg_scale "${CFG_SCALE}" --fps "${FPS}" \ + --output_path "${OUTPUT}" diff --git a/code/inference/dynamic_spatialvid/README.md b/code/inference/dynamic_spatialvid/README.md new file mode 100644 index 0000000000000000000000000000000000000000..0b8f205eed74408be51e16a9168da671cf4e2039 --- /dev/null +++ b/code/inference/dynamic_spatialvid/README.md @@ -0,0 +1,35 @@ +# Dynamic SpatialVID Inference + +Wrappers for checkpoints trained on the dynamic SpatialVID motion-filtered pool. They call `inference/unified_inference.py` and mirror the six public dynamic rows: + +| Script | `--memory_type` | Training recipe | +| --- | --- | --- | +| `run_infer_dyn_ctx1.sh` | `context_k1` | `train/dynamic_spatialvid/run_dyn_ctx1.sh` | +| `run_infer_dyn_ctx5.sh` | `context_k5` | `train/dynamic_spatialvid/run_dyn_ctx5.sh` | +| `run_infer_dyn_ctx20.sh` | `context_k20` | `train/dynamic_spatialvid/run_dyn_ctx20.sh` | +| `run_infer_dyn_spatial_mem.sh` | `spatial_mem` | `train/dynamic_spatialvid/run_dyn_spatial_mem.sh` | +| `run_infer_dyn_block_wise_ssm.sh` | `block_wise_ssm` | `train/dynamic_spatialvid/run_dyn_block_wise_ssm.sh` | +| `run_infer_dyn_videossm_hybrid.sh` | `videossm_hybrid` | `train/dynamic_spatialvid/run_dyn_videossm_hybrid.sh` | + +```bash +export WAN_BASE_MODEL=/path/to/Wan2.1-T2V-1.3B +export CKPT=/path/to/retrained_dynamic_spatial_mem/epoch-0.safetensors +PROMPT="A dynamic outdoor scene with a smooth camera move" \ + bash inference/dynamic_spatialvid/run_infer_dyn_spatial_mem.sh +``` + +To run all six rows: + +```bash +CKPT_DIR=./ckpts/dynamic_spatialvid bash inference/dynamic_spatialvid/run_infer_all_dyn.sh +``` + +## Demo Selection + +Dynamic demos should be selected from training scenes by random replay, then manually picked: + +1. Sample candidate rows from `metadata_train.csv` or a small public subset such as `metadata_train_sample.csv`. +2. Replay the same scene/prompt/action with all six dynamic checkpoints. +3. Pick representative successes for the public preview grid. + +Evaluation for the dynamic benchmark is intentionally TODO for now; only training and inference wrappers are public. diff --git a/code/inference/dynamic_spatialvid/common_env.sh b/code/inference/dynamic_spatialvid/common_env.sh new file mode 100644 index 0000000000000000000000000000000000000000..67f7c156ce7ac8b9e89617c44ecf8b1bc5eda86c --- /dev/null +++ b/code/inference/dynamic_spatialvid/common_env.sh @@ -0,0 +1,9 @@ +#!/bin/bash +# Shared dynamic SpatialVID inference environment. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/../_shared/common_env_infer.sh" + +INFER_OUTPUT_ROOT="${INFER_OUTPUT_ROOT:-${REPO_ROOT}/inference_outputs/dynamic_spatialvid}" +mkdir -p "${INFER_OUTPUT_ROOT}" +export INFER_OUTPUT_ROOT diff --git a/code/inference/dynamic_spatialvid/run_infer_all_dyn.sh b/code/inference/dynamic_spatialvid/run_infer_all_dyn.sh new file mode 100644 index 0000000000000000000000000000000000000000..760c890a80f5292f02f15529344aa2f460378f20 --- /dev/null +++ b/code/inference/dynamic_spatialvid/run_infer_all_dyn.sh @@ -0,0 +1,18 @@ +#!/bin/bash +# Run all dynamic SpatialVID inference rows from CKPT_DIR. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +: "${CKPT_DIR:?ERROR: set CKPT_DIR to a directory with dynamic row checkpoints}" + +run_one() { + local row="$1" + local script="$2" + CKPT="${CKPT_DIR}/${row}/epoch-0.safetensors" bash "${SCRIPT_DIR}/${script}" +} + +run_one ctx1 run_infer_dyn_ctx1.sh +run_one ctx5 run_infer_dyn_ctx5.sh +run_one ctx20 run_infer_dyn_ctx20.sh +run_one spatial_mem run_infer_dyn_spatial_mem.sh +run_one block_wise_ssm run_infer_dyn_block_wise_ssm.sh +run_one videossm_hybrid run_infer_dyn_videossm_hybrid.sh diff --git a/code/inference/dynamic_spatialvid/run_infer_dyn_block_wise_ssm.sh b/code/inference/dynamic_spatialvid/run_infer_dyn_block_wise_ssm.sh new file mode 100644 index 0000000000000000000000000000000000000000..60bffb285df3ca538f441496851ec54fe0576848 --- /dev/null +++ b/code/inference/dynamic_spatialvid/run_infer_dyn_block_wise_ssm.sh @@ -0,0 +1,19 @@ +#!/bin/bash +# Dynamic SpatialVID inference: Block-wise SSM. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +: "${CKPT:?ERROR: set CKPT to your dynamic block_wise_ssm checkpoint path}" +OUTPUT="${INFER_OUTPUT_ROOT}/dyn_block_wise_ssm_$(date +%Y%m%d_%H%M%S).mp4" +EXTRA_ARGS=() +[ -n "${CONTEXT_IMAGE}" ] && EXTRA_ARGS+=(--context_image "${CONTEXT_IMAGE}") +[ -n "${ACTION_PATH}" ] && EXTRA_ARGS+=(--action_path "${ACTION_PATH}") +[ -n "${TOKENIZER_PATH}" ] && EXTRA_ARGS+=(--tokenizer_path "${TOKENIZER_PATH}") + +python inference/unified_inference.py \ + --ckpt "${CKPT}" --memory_type block_wise_ssm --base_model "${WAN_BASE_MODEL}" \ + --prompt "${PROMPT}" "${EXTRA_ARGS[@]}" \ + --height "${HEIGHT}" --width "${WIDTH}" --num_frames "${NUM_FRAMES}" \ + --seed "${SEED}" --num_inference_steps "${NUM_INFERENCE_STEPS}" \ + --sigma_shift "${SIGMA_SHIFT}" --cfg_scale "${CFG_SCALE}" --fps "${FPS}" \ + --output_path "${OUTPUT}" diff --git a/code/inference/dynamic_spatialvid/run_infer_dyn_ctx1.sh b/code/inference/dynamic_spatialvid/run_infer_dyn_ctx1.sh new file mode 100644 index 0000000000000000000000000000000000000000..efda71ac4a969fb734616cb95038f32b9ba6465b --- /dev/null +++ b/code/inference/dynamic_spatialvid/run_infer_dyn_ctx1.sh @@ -0,0 +1,19 @@ +#!/bin/bash +# Dynamic SpatialVID inference: Context K=1. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +: "${CKPT:?ERROR: set CKPT to your dynamic ctx1 checkpoint path}" +OUTPUT="${INFER_OUTPUT_ROOT}/dyn_ctx1_$(date +%Y%m%d_%H%M%S).mp4" +EXTRA_ARGS=() +[ -n "${CONTEXT_IMAGE}" ] && EXTRA_ARGS+=(--context_image "${CONTEXT_IMAGE}") +[ -n "${ACTION_PATH}" ] && EXTRA_ARGS+=(--action_path "${ACTION_PATH}") +[ -n "${TOKENIZER_PATH}" ] && EXTRA_ARGS+=(--tokenizer_path "${TOKENIZER_PATH}") + +python inference/unified_inference.py \ + --ckpt "${CKPT}" --memory_type context_k1 --base_model "${WAN_BASE_MODEL}" \ + --prompt "${PROMPT}" "${EXTRA_ARGS[@]}" \ + --height "${HEIGHT}" --width "${WIDTH}" --num_frames "${NUM_FRAMES}" \ + --seed "${SEED}" --num_inference_steps "${NUM_INFERENCE_STEPS}" \ + --sigma_shift "${SIGMA_SHIFT:-5.0}" --cfg_scale "${CFG_SCALE}" --fps "${FPS}" \ + --output_path "${OUTPUT}" diff --git a/code/inference/dynamic_spatialvid/run_infer_dyn_ctx20.sh b/code/inference/dynamic_spatialvid/run_infer_dyn_ctx20.sh new file mode 100644 index 0000000000000000000000000000000000000000..17491bef192e18fa08d281bf5117d5a4a098eaa2 --- /dev/null +++ b/code/inference/dynamic_spatialvid/run_infer_dyn_ctx20.sh @@ -0,0 +1,19 @@ +#!/bin/bash +# Dynamic SpatialVID inference: Context K=20. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +: "${CKPT:?ERROR: set CKPT to your dynamic ctx20 checkpoint path}" +OUTPUT="${INFER_OUTPUT_ROOT}/dyn_ctx20_$(date +%Y%m%d_%H%M%S).mp4" +EXTRA_ARGS=() +[ -n "${CONTEXT_IMAGE}" ] && EXTRA_ARGS+=(--context_image "${CONTEXT_IMAGE}") +[ -n "${ACTION_PATH}" ] && EXTRA_ARGS+=(--action_path "${ACTION_PATH}") +[ -n "${TOKENIZER_PATH}" ] && EXTRA_ARGS+=(--tokenizer_path "${TOKENIZER_PATH}") + +python inference/unified_inference.py \ + --ckpt "${CKPT}" --memory_type context_k20 --base_model "${WAN_BASE_MODEL}" \ + --prompt "${PROMPT}" "${EXTRA_ARGS[@]}" \ + --height "${HEIGHT}" --width "${WIDTH}" --num_frames "${NUM_FRAMES}" \ + --seed "${SEED}" --num_inference_steps "${NUM_INFERENCE_STEPS}" \ + --sigma_shift "${SIGMA_SHIFT:-5.0}" --cfg_scale "${CFG_SCALE}" --fps "${FPS}" \ + --output_path "${OUTPUT}" diff --git a/code/inference/dynamic_spatialvid/run_infer_dyn_ctx5.sh b/code/inference/dynamic_spatialvid/run_infer_dyn_ctx5.sh new file mode 100644 index 0000000000000000000000000000000000000000..a91bd59c60542f04d73866a868cfd60572893743 --- /dev/null +++ b/code/inference/dynamic_spatialvid/run_infer_dyn_ctx5.sh @@ -0,0 +1,19 @@ +#!/bin/bash +# Dynamic SpatialVID inference: Context K=5. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +: "${CKPT:?ERROR: set CKPT to your dynamic ctx5 checkpoint path}" +OUTPUT="${INFER_OUTPUT_ROOT}/dyn_ctx5_$(date +%Y%m%d_%H%M%S).mp4" +EXTRA_ARGS=() +[ -n "${CONTEXT_IMAGE}" ] && EXTRA_ARGS+=(--context_image "${CONTEXT_IMAGE}") +[ -n "${ACTION_PATH}" ] && EXTRA_ARGS+=(--action_path "${ACTION_PATH}") +[ -n "${TOKENIZER_PATH}" ] && EXTRA_ARGS+=(--tokenizer_path "${TOKENIZER_PATH}") + +python inference/unified_inference.py \ + --ckpt "${CKPT}" --memory_type context_k5 --base_model "${WAN_BASE_MODEL}" \ + --prompt "${PROMPT}" "${EXTRA_ARGS[@]}" \ + --height "${HEIGHT}" --width "${WIDTH}" --num_frames "${NUM_FRAMES}" \ + --seed "${SEED}" --num_inference_steps "${NUM_INFERENCE_STEPS}" \ + --sigma_shift "${SIGMA_SHIFT:-5.0}" --cfg_scale "${CFG_SCALE}" --fps "${FPS}" \ + --output_path "${OUTPUT}" diff --git a/code/inference/dynamic_spatialvid/run_infer_dyn_spatial_mem.sh b/code/inference/dynamic_spatialvid/run_infer_dyn_spatial_mem.sh new file mode 100644 index 0000000000000000000000000000000000000000..90d527b2b297de3fafa9f80bf11a3af9b2e336c8 --- /dev/null +++ b/code/inference/dynamic_spatialvid/run_infer_dyn_spatial_mem.sh @@ -0,0 +1,19 @@ +#!/bin/bash +# Dynamic SpatialVID inference: Spatial Memory. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +: "${CKPT:?ERROR: set CKPT to your dynamic spatial_mem checkpoint path}" +OUTPUT="${INFER_OUTPUT_ROOT}/dyn_spatial_mem_$(date +%Y%m%d_%H%M%S).mp4" +EXTRA_ARGS=() +[ -n "${CONTEXT_IMAGE}" ] && EXTRA_ARGS+=(--context_image "${CONTEXT_IMAGE}") +[ -n "${ACTION_PATH}" ] && EXTRA_ARGS+=(--action_path "${ACTION_PATH}") +[ -n "${TOKENIZER_PATH}" ] && EXTRA_ARGS+=(--tokenizer_path "${TOKENIZER_PATH}") + +python inference/unified_inference.py \ + --ckpt "${CKPT}" --memory_type spatial_mem --base_model "${WAN_BASE_MODEL}" \ + --prompt "${PROMPT}" "${EXTRA_ARGS[@]}" \ + --height "${HEIGHT}" --width "${WIDTH}" --num_frames "${NUM_FRAMES}" \ + --seed "${SEED}" --num_inference_steps "${NUM_INFERENCE_STEPS}" \ + --sigma_shift "${SIGMA_SHIFT}" --cfg_scale "${CFG_SCALE}" --fps "${FPS}" \ + --output_path "${OUTPUT}" diff --git a/code/inference/dynamic_spatialvid/run_infer_dyn_videossm_hybrid.sh b/code/inference/dynamic_spatialvid/run_infer_dyn_videossm_hybrid.sh new file mode 100644 index 0000000000000000000000000000000000000000..82f3aa189acfcf06f5ac1e4969023fb716553629 --- /dev/null +++ b/code/inference/dynamic_spatialvid/run_infer_dyn_videossm_hybrid.sh @@ -0,0 +1,19 @@ +#!/bin/bash +# Dynamic SpatialVID inference: legacy VideoSSM hybrid. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +: "${CKPT:?ERROR: set CKPT to your dynamic videossm_hybrid checkpoint path}" +OUTPUT="${INFER_OUTPUT_ROOT}/dyn_videossm_hybrid_$(date +%Y%m%d_%H%M%S).mp4" +EXTRA_ARGS=() +[ -n "${CONTEXT_IMAGE}" ] && EXTRA_ARGS+=(--context_image "${CONTEXT_IMAGE}") +[ -n "${ACTION_PATH}" ] && EXTRA_ARGS+=(--action_path "${ACTION_PATH}") +[ -n "${TOKENIZER_PATH}" ] && EXTRA_ARGS+=(--tokenizer_path "${TOKENIZER_PATH}") + +python inference/unified_inference.py \ + --ckpt "${CKPT}" --memory_type videossm_hybrid --base_model "${WAN_BASE_MODEL}" \ + --prompt "${PROMPT}" "${EXTRA_ARGS[@]}" \ + --height "${HEIGHT}" --width "${WIDTH}" --num_frames "${NUM_FRAMES}" \ + --seed "${SEED}" --num_inference_steps "${NUM_INFERENCE_STEPS}" \ + --sigma_shift "${SIGMA_SHIFT}" --cfg_scale "${CFG_SCALE}" --fps "${FPS}" \ + --output_path "${OUTPUT}" diff --git a/code/inference/memory_baselines_basic/common_env.sh b/code/inference/memory_baselines_basic/common_env.sh new file mode 100644 index 0000000000000000000000000000000000000000..22c0665846e16861293052698db694669c8119d6 --- /dev/null +++ b/code/inference/memory_baselines_basic/common_env.sh @@ -0,0 +1,4 @@ +#!/bin/bash +# Common inference wrapper for memory_baselines_basic. +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/../_shared/common_env_infer.sh" diff --git a/code/inference/memory_baselines_basic/run_infer_all.sh b/code/inference/memory_baselines_basic/run_infer_all.sh new file mode 100644 index 0000000000000000000000000000000000000000..cfea2793fd815d23e0bacdb41edbd93d5414239a --- /dev/null +++ b/code/inference/memory_baselines_basic/run_infer_all.sh @@ -0,0 +1,47 @@ +#!/bin/bash +# Run all memory baseline inferences sequentially. +# Requires: CKPT_DIR pointing to a directory with per-row checkpoint folders +# (same layout as HF repo: context_k1/epoch-0.safetensors, etc.) +# +# Usage: +# CKPT_DIR=./ckpts bash inference/memory_baselines_basic/run_infer_all.sh +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +: "${CKPT_DIR:?ERROR: set CKPT_DIR to your checkpoints directory}" + +declare -A CKPT_MAP=( + [run_infer_no_memory.sh]="no_memory_extra_two_chunk" + [run_infer_framepack_weight.sh]="framepack_weight_only" + [run_infer_framepack_len_r2.sh]="framepack_lencompress_r2" + [run_infer_framepack_len_r4.sh]="framepack_lencompress_r4" + [run_infer_framepack_hybrid_r2.sh]="framepack_hybrid_r2_weight_two_chunk" + [run_infer_framepack_hybrid_r4.sh]="framepack_hybrid_r4_weight_two_chunk" + [run_infer_spatial_mem.sh]="spatial_mem" + [run_infer_spatial_concat_text.sh]="spatial_concat_text_two_chunk" + [run_infer_spatial_inject_none.sh]="spatial_inject_none_two_chunk" + [run_infer_spatial_cross_attn_readout.sh]="spatial_cross_attn_readout_two_chunk" + [run_infer_videossm_hybrid.sh]="videossm_hybrid" + [run_infer_block_wise_ssm.sh]="block_wise_ssm_two_chunk" +) + +for script in "${SCRIPT_DIR}"/run_infer_*.sh; do + name="$(basename "${script}")" + [[ "${name}" == "run_infer_all.sh" ]] && continue + ckpt_folder="${CKPT_MAP[${name}]:-}" + if [ -z "${ckpt_folder}" ]; then + echo "[run_infer_all] SKIP: no checkpoint mapping for ${name}" + continue + fi + ckpt_file="${CKPT_DIR}/${ckpt_folder}/epoch-0.safetensors" + if [ ! -f "${ckpt_file}" ]; then + echo "[run_infer_all] SKIP: ${ckpt_file} not found" + continue + fi + echo "========================================" + echo "[run_infer_all] Running ${name} with ${ckpt_file}" + echo "========================================" + CKPT="${ckpt_file}" bash "${script}" +done + +echo "[run_infer_all] Done." diff --git a/code/inference/memory_baselines_basic/run_infer_block_wise_ssm.sh b/code/inference/memory_baselines_basic/run_infer_block_wise_ssm.sh new file mode 100644 index 0000000000000000000000000000000000000000..f260c060810ef7dcc48edbc90269269e2d19b3d7 --- /dev/null +++ b/code/inference/memory_baselines_basic/run_infer_block_wise_ssm.sh @@ -0,0 +1,28 @@ +#!/bin/bash +# Inference: block_wise_ssm +# Corresponds to: train/memory_baselines_basic/run_ablation_block_wise_ssm_two_chunk.sh +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" + +: "${CKPT:?ERROR: set CKPT to your block_wise_ssm checkpoint path}" +OUTPUT="${INFER_OUTPUT_ROOT}/block_wise_ssm_$(date +%Y%m%d_%H%M%S).mp4" + +EXTRA_ARGS=() +if [ -n "${CONTEXT_IMAGE}" ]; then + EXTRA_ARGS+=(--context_image "${CONTEXT_IMAGE}") +fi +if [ -n "${ACTION_PATH}" ]; then + EXTRA_ARGS+=(--action_path "${ACTION_PATH}") +fi + +python inference/unified_inference.py \ + --ckpt "${CKPT}" \ + --memory_type block_wise_ssm \ + --base_model "${WAN_BASE_MODEL}" \ + --prompt "${PROMPT}" \ + "${EXTRA_ARGS[@]}" \ + --height "${HEIGHT}" --width "${WIDTH}" --num_frames "${NUM_FRAMES}" \ + --seed "${SEED}" --num_inference_steps "${NUM_INFERENCE_STEPS}" \ + --sigma_shift "${SIGMA_SHIFT}" --cfg_scale "${CFG_SCALE}" --fps "${FPS}" \ + --output_path "${OUTPUT}" diff --git a/code/inference/memory_baselines_basic/run_infer_framepack_hybrid_r2.sh b/code/inference/memory_baselines_basic/run_infer_framepack_hybrid_r2.sh new file mode 100644 index 0000000000000000000000000000000000000000..398be8ff75cabc057756eb9d014c1ad0c1152484 --- /dev/null +++ b/code/inference/memory_baselines_basic/run_infer_framepack_hybrid_r2.sh @@ -0,0 +1,28 @@ +#!/bin/bash +# Inference: framepack_hybrid_r2 +# Corresponds to: train/memory_baselines_basic/run_ablation_framepack_hybrid_r2_weight_two_chunk.sh +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" + +: "${CKPT:?ERROR: set CKPT to your framepack_hybrid_r2 checkpoint path}" +OUTPUT="${INFER_OUTPUT_ROOT}/framepack_hybrid_r2_$(date +%Y%m%d_%H%M%S).mp4" + +EXTRA_ARGS=() +if [ -n "${CONTEXT_IMAGE}" ]; then + EXTRA_ARGS+=(--context_image "${CONTEXT_IMAGE}") +fi +if [ -n "${ACTION_PATH}" ]; then + EXTRA_ARGS+=(--action_path "${ACTION_PATH}") +fi + +python inference/unified_inference.py \ + --ckpt "${CKPT}" \ + --memory_type framepack_hybrid_r2 \ + --base_model "${WAN_BASE_MODEL}" \ + --prompt "${PROMPT}" \ + "${EXTRA_ARGS[@]}" \ + --height "${HEIGHT}" --width "${WIDTH}" --num_frames "${NUM_FRAMES}" \ + --seed "${SEED}" --num_inference_steps "${NUM_INFERENCE_STEPS}" \ + --sigma_shift "${SIGMA_SHIFT}" --cfg_scale "${CFG_SCALE}" --fps "${FPS}" \ + --output_path "${OUTPUT}" diff --git a/code/inference/memory_baselines_basic/run_infer_framepack_hybrid_r4.sh b/code/inference/memory_baselines_basic/run_infer_framepack_hybrid_r4.sh new file mode 100644 index 0000000000000000000000000000000000000000..8203138c71f3936288a94a49e65ecde0f1161c5d --- /dev/null +++ b/code/inference/memory_baselines_basic/run_infer_framepack_hybrid_r4.sh @@ -0,0 +1,28 @@ +#!/bin/bash +# Inference: framepack_hybrid_r4 +# Corresponds to: train/memory_baselines_basic/run_ablation_framepack_hybrid_r4_weight_two_chunk.sh +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" + +: "${CKPT:?ERROR: set CKPT to your framepack_hybrid_r4 checkpoint path}" +OUTPUT="${INFER_OUTPUT_ROOT}/framepack_hybrid_r4_$(date +%Y%m%d_%H%M%S).mp4" + +EXTRA_ARGS=() +if [ -n "${CONTEXT_IMAGE}" ]; then + EXTRA_ARGS+=(--context_image "${CONTEXT_IMAGE}") +fi +if [ -n "${ACTION_PATH}" ]; then + EXTRA_ARGS+=(--action_path "${ACTION_PATH}") +fi + +python inference/unified_inference.py \ + --ckpt "${CKPT}" \ + --memory_type framepack_hybrid_r4 \ + --base_model "${WAN_BASE_MODEL}" \ + --prompt "${PROMPT}" \ + "${EXTRA_ARGS[@]}" \ + --height "${HEIGHT}" --width "${WIDTH}" --num_frames "${NUM_FRAMES}" \ + --seed "${SEED}" --num_inference_steps "${NUM_INFERENCE_STEPS}" \ + --sigma_shift "${SIGMA_SHIFT}" --cfg_scale "${CFG_SCALE}" --fps "${FPS}" \ + --output_path "${OUTPUT}" diff --git a/code/inference/memory_baselines_basic/run_infer_framepack_len_r2.sh b/code/inference/memory_baselines_basic/run_infer_framepack_len_r2.sh new file mode 100644 index 0000000000000000000000000000000000000000..9473c3f5d629e7a30faa26b4459b410134bff74c --- /dev/null +++ b/code/inference/memory_baselines_basic/run_infer_framepack_len_r2.sh @@ -0,0 +1,28 @@ +#!/bin/bash +# Inference: framepack_len_r2 +# Corresponds to: train/memory_baselines_basic/run_ablation_framepack_len_r2_two_chunk.sh +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" + +: "${CKPT:?ERROR: set CKPT to your framepack_len_r2 checkpoint path}" +OUTPUT="${INFER_OUTPUT_ROOT}/framepack_len_r2_$(date +%Y%m%d_%H%M%S).mp4" + +EXTRA_ARGS=() +if [ -n "${CONTEXT_IMAGE}" ]; then + EXTRA_ARGS+=(--context_image "${CONTEXT_IMAGE}") +fi +if [ -n "${ACTION_PATH}" ]; then + EXTRA_ARGS+=(--action_path "${ACTION_PATH}") +fi + +python inference/unified_inference.py \ + --ckpt "${CKPT}" \ + --memory_type framepack_len_r2 \ + --base_model "${WAN_BASE_MODEL}" \ + --prompt "${PROMPT}" \ + "${EXTRA_ARGS[@]}" \ + --height "${HEIGHT}" --width "${WIDTH}" --num_frames "${NUM_FRAMES}" \ + --seed "${SEED}" --num_inference_steps "${NUM_INFERENCE_STEPS}" \ + --sigma_shift "${SIGMA_SHIFT}" --cfg_scale "${CFG_SCALE}" --fps "${FPS}" \ + --output_path "${OUTPUT}" diff --git a/code/inference/memory_baselines_basic/run_infer_framepack_len_r4.sh b/code/inference/memory_baselines_basic/run_infer_framepack_len_r4.sh new file mode 100644 index 0000000000000000000000000000000000000000..8de56d9705c2264947eb00abad37d563a91a71f7 --- /dev/null +++ b/code/inference/memory_baselines_basic/run_infer_framepack_len_r4.sh @@ -0,0 +1,28 @@ +#!/bin/bash +# Inference: framepack_len_r4 +# Corresponds to: train/memory_baselines_basic/run_ablation_framepack_len_r4_two_chunk.sh +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" + +: "${CKPT:?ERROR: set CKPT to your framepack_len_r4 checkpoint path}" +OUTPUT="${INFER_OUTPUT_ROOT}/framepack_len_r4_$(date +%Y%m%d_%H%M%S).mp4" + +EXTRA_ARGS=() +if [ -n "${CONTEXT_IMAGE}" ]; then + EXTRA_ARGS+=(--context_image "${CONTEXT_IMAGE}") +fi +if [ -n "${ACTION_PATH}" ]; then + EXTRA_ARGS+=(--action_path "${ACTION_PATH}") +fi + +python inference/unified_inference.py \ + --ckpt "${CKPT}" \ + --memory_type framepack_len_r4 \ + --base_model "${WAN_BASE_MODEL}" \ + --prompt "${PROMPT}" \ + "${EXTRA_ARGS[@]}" \ + --height "${HEIGHT}" --width "${WIDTH}" --num_frames "${NUM_FRAMES}" \ + --seed "${SEED}" --num_inference_steps "${NUM_INFERENCE_STEPS}" \ + --sigma_shift "${SIGMA_SHIFT}" --cfg_scale "${CFG_SCALE}" --fps "${FPS}" \ + --output_path "${OUTPUT}" diff --git a/code/inference/memory_baselines_basic/run_infer_framepack_weight.sh b/code/inference/memory_baselines_basic/run_infer_framepack_weight.sh new file mode 100644 index 0000000000000000000000000000000000000000..eb36a7642bb725b1a8793a14b7cacfc136e8f805 --- /dev/null +++ b/code/inference/memory_baselines_basic/run_infer_framepack_weight.sh @@ -0,0 +1,28 @@ +#!/bin/bash +# Inference: framepack_weight +# Corresponds to: train/memory_baselines_basic/run_ablation_framepack_weight_two_chunk.sh +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" + +: "${CKPT:?ERROR: set CKPT to your framepack_weight checkpoint path}" +OUTPUT="${INFER_OUTPUT_ROOT}/framepack_weight_$(date +%Y%m%d_%H%M%S).mp4" + +EXTRA_ARGS=() +if [ -n "${CONTEXT_IMAGE}" ]; then + EXTRA_ARGS+=(--context_image "${CONTEXT_IMAGE}") +fi +if [ -n "${ACTION_PATH}" ]; then + EXTRA_ARGS+=(--action_path "${ACTION_PATH}") +fi + +python inference/unified_inference.py \ + --ckpt "${CKPT}" \ + --memory_type framepack_weight \ + --base_model "${WAN_BASE_MODEL}" \ + --prompt "${PROMPT}" \ + "${EXTRA_ARGS[@]}" \ + --height "${HEIGHT}" --width "${WIDTH}" --num_frames "${NUM_FRAMES}" \ + --seed "${SEED}" --num_inference_steps "${NUM_INFERENCE_STEPS}" \ + --sigma_shift "${SIGMA_SHIFT}" --cfg_scale "${CFG_SCALE}" --fps "${FPS}" \ + --output_path "${OUTPUT}" diff --git a/code/inference/memory_baselines_basic/run_infer_no_memory.sh b/code/inference/memory_baselines_basic/run_infer_no_memory.sh new file mode 100644 index 0000000000000000000000000000000000000000..60b6e4af54d52524a1e67445b600b45b7b02a1c2 --- /dev/null +++ b/code/inference/memory_baselines_basic/run_infer_no_memory.sh @@ -0,0 +1,28 @@ +#!/bin/bash +# Inference: No memory baseline (I2V floor). +# Corresponds to: train/memory_baselines_basic/run_ablation_no_memory_baseline_two_chunk.sh +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" + +: "${CKPT:?ERROR: set CKPT to your no-memory checkpoint path}" +OUTPUT="${INFER_OUTPUT_ROOT}/no_memory_$(date +%Y%m%d_%H%M%S).mp4" + +EXTRA_ARGS=() +if [ -n "${CONTEXT_IMAGE}" ]; then + EXTRA_ARGS+=(--context_image "${CONTEXT_IMAGE}") +fi +if [ -n "${ACTION_PATH}" ]; then + EXTRA_ARGS+=(--action_path "${ACTION_PATH}") +fi + +python inference/unified_inference.py \ + --ckpt "${CKPT}" \ + --memory_type no_memory \ + --base_model "${WAN_BASE_MODEL}" \ + --prompt "${PROMPT}" \ + "${EXTRA_ARGS[@]}" \ + --height "${HEIGHT}" --width "${WIDTH}" --num_frames "${NUM_FRAMES}" \ + --seed "${SEED}" --num_inference_steps "${NUM_INFERENCE_STEPS}" \ + --sigma_shift "${SIGMA_SHIFT}" --cfg_scale "${CFG_SCALE}" --fps "${FPS}" \ + --output_path "${OUTPUT}" diff --git a/code/inference/memory_baselines_basic/run_infer_spatial_concat_text.sh b/code/inference/memory_baselines_basic/run_infer_spatial_concat_text.sh new file mode 100644 index 0000000000000000000000000000000000000000..4ea056113e5caf7a5d6b8d2532255d709146320f --- /dev/null +++ b/code/inference/memory_baselines_basic/run_infer_spatial_concat_text.sh @@ -0,0 +1,28 @@ +#!/bin/bash +# Inference: spatial_concat_text +# Corresponds to: train/memory_baselines_basic/run_ablation_spatial_concat_text_two_chunk.sh +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" + +: "${CKPT:?ERROR: set CKPT to your spatial_concat_text checkpoint path}" +OUTPUT="${INFER_OUTPUT_ROOT}/spatial_concat_text_$(date +%Y%m%d_%H%M%S).mp4" + +EXTRA_ARGS=() +if [ -n "${CONTEXT_IMAGE}" ]; then + EXTRA_ARGS+=(--context_image "${CONTEXT_IMAGE}") +fi +if [ -n "${ACTION_PATH}" ]; then + EXTRA_ARGS+=(--action_path "${ACTION_PATH}") +fi + +python inference/unified_inference.py \ + --ckpt "${CKPT}" \ + --memory_type spatial_concat_text \ + --base_model "${WAN_BASE_MODEL}" \ + --prompt "${PROMPT}" \ + "${EXTRA_ARGS[@]}" \ + --height "${HEIGHT}" --width "${WIDTH}" --num_frames "${NUM_FRAMES}" \ + --seed "${SEED}" --num_inference_steps "${NUM_INFERENCE_STEPS}" \ + --sigma_shift "${SIGMA_SHIFT}" --cfg_scale "${CFG_SCALE}" --fps "${FPS}" \ + --output_path "${OUTPUT}" diff --git a/code/inference/memory_baselines_basic/run_infer_spatial_cross_attn_readout.sh b/code/inference/memory_baselines_basic/run_infer_spatial_cross_attn_readout.sh new file mode 100644 index 0000000000000000000000000000000000000000..8ce43dfcf5b7ff9e210acc1accf30ff064c3ba51 --- /dev/null +++ b/code/inference/memory_baselines_basic/run_infer_spatial_cross_attn_readout.sh @@ -0,0 +1,28 @@ +#!/bin/bash +# Inference: spatial_cross_attn_readout +# Corresponds to: train/memory_baselines_basic/run_ablation_spatial_cross_attn_readout_two_chunk.sh +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" + +: "${CKPT:?ERROR: set CKPT to your spatial_cross_attn_readout checkpoint path}" +OUTPUT="${INFER_OUTPUT_ROOT}/spatial_cross_attn_readout_$(date +%Y%m%d_%H%M%S).mp4" + +EXTRA_ARGS=() +if [ -n "${CONTEXT_IMAGE}" ]; then + EXTRA_ARGS+=(--context_image "${CONTEXT_IMAGE}") +fi +if [ -n "${ACTION_PATH}" ]; then + EXTRA_ARGS+=(--action_path "${ACTION_PATH}") +fi + +python inference/unified_inference.py \ + --ckpt "${CKPT}" \ + --memory_type spatial_cross_attn_readout \ + --base_model "${WAN_BASE_MODEL}" \ + --prompt "${PROMPT}" \ + "${EXTRA_ARGS[@]}" \ + --height "${HEIGHT}" --width "${WIDTH}" --num_frames "${NUM_FRAMES}" \ + --seed "${SEED}" --num_inference_steps "${NUM_INFERENCE_STEPS}" \ + --sigma_shift "${SIGMA_SHIFT}" --cfg_scale "${CFG_SCALE}" --fps "${FPS}" \ + --output_path "${OUTPUT}" diff --git a/code/inference/memory_baselines_basic/run_infer_spatial_inject_none.sh b/code/inference/memory_baselines_basic/run_infer_spatial_inject_none.sh new file mode 100644 index 0000000000000000000000000000000000000000..92ca05fccd676993a44d2204c58972df7710d29b --- /dev/null +++ b/code/inference/memory_baselines_basic/run_infer_spatial_inject_none.sh @@ -0,0 +1,28 @@ +#!/bin/bash +# Inference: spatial_inject_none +# Corresponds to: train/memory_baselines_basic/run_ablation_spatial_inject_none_two_chunk.sh +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" + +: "${CKPT:?ERROR: set CKPT to your spatial_inject_none checkpoint path}" +OUTPUT="${INFER_OUTPUT_ROOT}/spatial_inject_none_$(date +%Y%m%d_%H%M%S).mp4" + +EXTRA_ARGS=() +if [ -n "${CONTEXT_IMAGE}" ]; then + EXTRA_ARGS+=(--context_image "${CONTEXT_IMAGE}") +fi +if [ -n "${ACTION_PATH}" ]; then + EXTRA_ARGS+=(--action_path "${ACTION_PATH}") +fi + +python inference/unified_inference.py \ + --ckpt "${CKPT}" \ + --memory_type spatial_inject_none \ + --base_model "${WAN_BASE_MODEL}" \ + --prompt "${PROMPT}" \ + "${EXTRA_ARGS[@]}" \ + --height "${HEIGHT}" --width "${WIDTH}" --num_frames "${NUM_FRAMES}" \ + --seed "${SEED}" --num_inference_steps "${NUM_INFERENCE_STEPS}" \ + --sigma_shift "${SIGMA_SHIFT}" --cfg_scale "${CFG_SCALE}" --fps "${FPS}" \ + --output_path "${OUTPUT}" diff --git a/code/inference/memory_baselines_basic/run_infer_spatial_mem.sh b/code/inference/memory_baselines_basic/run_infer_spatial_mem.sh new file mode 100644 index 0000000000000000000000000000000000000000..c48485daa1ce3fca92bd7ec4d941819bdbf82f14 --- /dev/null +++ b/code/inference/memory_baselines_basic/run_infer_spatial_mem.sh @@ -0,0 +1,28 @@ +#!/bin/bash +# Inference: spatial_mem +# Corresponds to: train/memory_baselines_basic/run_spatial_memory_baseline.sh +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" + +: "${CKPT:?ERROR: set CKPT to your spatial_mem checkpoint path}" +OUTPUT="${INFER_OUTPUT_ROOT}/spatial_mem_$(date +%Y%m%d_%H%M%S).mp4" + +EXTRA_ARGS=() +if [ -n "${CONTEXT_IMAGE}" ]; then + EXTRA_ARGS+=(--context_image "${CONTEXT_IMAGE}") +fi +if [ -n "${ACTION_PATH}" ]; then + EXTRA_ARGS+=(--action_path "${ACTION_PATH}") +fi + +python inference/unified_inference.py \ + --ckpt "${CKPT}" \ + --memory_type spatial_mem \ + --base_model "${WAN_BASE_MODEL}" \ + --prompt "${PROMPT}" \ + "${EXTRA_ARGS[@]}" \ + --height "${HEIGHT}" --width "${WIDTH}" --num_frames "${NUM_FRAMES}" \ + --seed "${SEED}" --num_inference_steps "${NUM_INFERENCE_STEPS}" \ + --sigma_shift "${SIGMA_SHIFT}" --cfg_scale "${CFG_SCALE}" --fps "${FPS}" \ + --output_path "${OUTPUT}" diff --git a/code/inference/memory_baselines_basic/run_infer_videossm_hybrid.sh b/code/inference/memory_baselines_basic/run_infer_videossm_hybrid.sh new file mode 100644 index 0000000000000000000000000000000000000000..e4fab4c404ccab2775f4e8785f285628d6ec2161 --- /dev/null +++ b/code/inference/memory_baselines_basic/run_infer_videossm_hybrid.sh @@ -0,0 +1,28 @@ +#!/bin/bash +# Inference: videossm_hybrid +# Corresponds to: train/memory_baselines_basic/run_videossm_hybrid_baseline.sh +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" + +: "${CKPT:?ERROR: set CKPT to your videossm_hybrid checkpoint path}" +OUTPUT="${INFER_OUTPUT_ROOT}/videossm_hybrid_$(date +%Y%m%d_%H%M%S).mp4" + +EXTRA_ARGS=() +if [ -n "${CONTEXT_IMAGE}" ]; then + EXTRA_ARGS+=(--context_image "${CONTEXT_IMAGE}") +fi +if [ -n "${ACTION_PATH}" ]; then + EXTRA_ARGS+=(--action_path "${ACTION_PATH}") +fi + +python inference/unified_inference.py \ + --ckpt "${CKPT}" \ + --memory_type videossm_hybrid \ + --base_model "${WAN_BASE_MODEL}" \ + --prompt "${PROMPT}" \ + "${EXTRA_ARGS[@]}" \ + --height "${HEIGHT}" --width "${WIDTH}" --num_frames "${NUM_FRAMES}" \ + --seed "${SEED}" --num_inference_steps "${NUM_INFERENCE_STEPS}" \ + --sigma_shift "${SIGMA_SHIFT}" --cfg_scale "${CFG_SCALE}" --fps "${FPS}" \ + --output_path "${OUTPUT}" diff --git a/code/inference/unified_inference.py b/code/inference/unified_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..3320da1a0ae982b16a302bc85627d089c73ca41b --- /dev/null +++ b/code/inference/unified_inference.py @@ -0,0 +1,390 @@ +#!/usr/bin/env python3 +""" +Unified single-chunk inference for all Echo-Memory memory families. + +Supports: no_memory, context_k1/k5/k20, framepack_weight, framepack_len_r2/r4, +framepack_hybrid_r2/r4, spatial_mem, spatial_concat_text, spatial_inject_none, +spatial_cross_attn_readout, videossm_hybrid, block_wise_ssm. + +Memory type can be specified explicitly via --memory_type or auto-detected +from the checkpoint path (--memory_type auto). + +Examples: + + # Auto-detect memory type from checkpoint path + python inference/unified_inference.py \ + --ckpt ./ckpts/context_k1/epoch-0.safetensors \ + --prompt "A toy bear on a table, the camera rotates around it" \ + --output_path output.mp4 + + # Explicit memory type + python inference/unified_inference.py \ + --ckpt ./ckpts/my_checkpoint.safetensors \ + --memory_type context_k1 \ + --prompt "A scene" \ + --output_path output.mp4 + + # With context image (first frame conditioning) + python inference/unified_inference.py \ + --ckpt ./ckpts/context_k1/epoch-0.safetensors \ + --context_image assets/opendomain_revisit/1774363417.png \ + --action_path env/action_rotation_left_45.json \ + --prompt "A toy bear on a table" \ + --output_path output.mp4 +""" +from __future__ import annotations + +import argparse +import os +import sys + +# Ensure repo root is in sys.path +_script_dir = os.path.dirname(os.path.abspath(__file__)) +_repo_root = os.path.abspath(os.path.join(_script_dir, "..")) +if _repo_root not in sys.path: + sys.path.insert(0, _repo_root) + +# memory_baseline_runtime has no heavy deps — safe to import at module level +from env.memory_baseline_runtime import ( + MemoryProfile, + MEMORY_PROFILE_REGISTRY, + infer_memory_profile_spec, +) + + +# --------------------------------------------------------------------------- +# Memory type → profile mapping +# --------------------------------------------------------------------------- + +# Friendly name → registry profile_id +_REGISTRY_ALIAS = { + "no_memory": "no_memory_extra_two_chunk", + "framepack_weight": "framepack_weight_only", + "framepack_len_r2": "framepack_lencompress_r2", + "framepack_len_r4": "framepack_lencompress_r4", + "framepack_hybrid_r2": "framepack_hybrid_r2_weight_two_chunk", + "framepack_hybrid_r4": "framepack_hybrid_r4_weight_two_chunk", + "spatial_mem": "spatial_mem", + "spatial_concat_text": "spatial_concat_text_two_chunk", + "spatial_inject_none": "spatial_inject_none_two_chunk", + "spatial_cross_attn_readout": "spatial_cross_attn_readout_two_chunk", + "geometry_spatial_mem": "geometry_spatial_mem", + "videossm_hybrid": "videossm_hybrid_legacy", + "block_wise_ssm": "block_wise_ssm_two_chunk", + "cgla_memory": "cgla_memory_two_chunk", + "prope_memory": "prope_memory_two_chunk", + "ucpe_memory": "ucpe_memory_two_chunk", +} + +# context_k* are not in the registry; they use default pipe flags +# with only context_override differing +_CONTEXT_K_PROFILES = { + "context_k1": MemoryProfile(context_override=1), + "context_k5": MemoryProfile(context_override=5), + "context_k20": MemoryProfile(context_override=20), +} + +# Build profile_id → MemoryProfile lookup from the registry +_REGISTRY_PROFILES = {spec.profile_id: spec.profile for spec in MEMORY_PROFILE_REGISTRY} + +ALL_MEMORY_TYPES = ["auto"] + sorted( + set(_REGISTRY_ALIAS.keys()) | set(_CONTEXT_K_PROFILES.keys()) +) + + +def resolve_memory_profile(memory_type: str, ckpt_path: str) -> MemoryProfile: + """Resolve --memory_type to a MemoryProfile.""" + if memory_type == "auto": + spec = infer_memory_profile_spec(ckpt_path) + if spec is None: + print( + f"[unified_inference] WARNING: --memory_type=auto but checkpoint path " + f"does not match any known memory profile. Running with no memory flags.\n" + f" ckpt: {ckpt_path}\n" + f" Hint: use --memory_type to specify explicitly.", + file=sys.stderr, + flush=True, + ) + return MemoryProfile() + print(f"[unified_inference] Auto-detected memory profile: {spec.profile_id}") + return spec.profile + + if memory_type in _CONTEXT_K_PROFILES: + print(f"[unified_inference] Using context learning profile: {memory_type}") + return _CONTEXT_K_PROFILES[memory_type] + + if memory_type in _REGISTRY_ALIAS: + profile_id = _REGISTRY_ALIAS[memory_type] + profile = _REGISTRY_PROFILES[profile_id] + print(f"[unified_inference] Using memory profile: {memory_type} ({profile_id})") + return profile + + print( + f"[unified_inference] ERROR: unknown --memory_type '{memory_type}'. " + f"Available: {', '.join(ALL_MEMORY_TYPES)}", + file=sys.stderr, + ) + sys.exit(1) + + +def apply_profile_to_pipe(pipe, profile: MemoryProfile) -> None: + """Apply a MemoryProfile directly to the pipeline object.""" + pipe.use_framepack_memory = bool(profile.use_framepack_memory) + pipe.context_temporal_decay = float(profile.context_temporal_decay or 1.0) + pipe.context_attention_weight = float(profile.context_attention_weight or 1.0) + pipe.use_framepack_length_compress = bool(profile.use_framepack_length_compress) + pipe.framepack_ratio = int(profile.framepack_ratio or 2) + pipe.use_spatial_memory = bool(profile.use_spatial_memory) + pipe.spatial_memory_tokens = int(profile.spatial_memory_tokens or 64) + if profile.spatial_memory_inject_mode: + pipe.spatial_memory_inject_mode = str(profile.spatial_memory_inject_mode) + pipe.use_spatial_memory_legacy = bool(profile.use_spatial_memory_legacy) + pipe.use_geometry_spatial_memory = bool(profile.use_geometry_spatial_memory) + if profile.geometry_spatial_memory_inject_mode: + pipe.geometry_spatial_memory_inject_mode = str( + profile.geometry_spatial_memory_inject_mode + ) + pipe.use_block_wise_ssm = bool(getattr(profile, "use_block_wise_ssm", False)) + pipe.use_videossm_hybrid = bool(getattr(profile, "use_videossm_hybrid", False)) + pipe.use_cgla_memory = bool(getattr(profile, "use_cgla_memory", False)) + # Warn if spatial memory requested but module not loaded from checkpoint + if ( + pipe.use_spatial_memory + and not pipe.use_spatial_memory_legacy + and getattr(pipe, "spatial_memory_module", None) is None + ): + raise RuntimeError( + "Spatial token-grid profile requested, but checkpoint has no " + "spatial_memory_module weights. Refusing to silently substitute the " + "legacy adaptive pool; select an explicit legacy profile instead." + ) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Unified single-chunk inference for all Echo-Memory memory families.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Memory types: + auto Auto-detect from checkpoint path + no_memory No memory (I2V floor baseline) + context_k1/k5/k20 Raw context with 1/5/20 frames + framepack_weight FramePack temporal decay reweighting + framepack_len_r2/r4 FramePack length compression ratio 2/4 + framepack_hybrid_r2/r4 FramePack hybrid (length + weight) + spatial_mem Spatial grid memory (64 tokens) + spatial_concat_text Spatial memory via text KV concatenation + spatial_inject_none Spatial memory with withheld read-out + spatial_cross_attn_readout Spatial memory via cross-attention + videossm_hybrid Legacy VideoSSM hybrid (temporal-conv baseline) + block_wise_ssm Block-wise recurrent SSM (paper-aligned) + cgla_memory Camera-guided linear attention (pose-conditioned GLA) + prope_memory PRoPE: CGLA + camera-pose rotary PE (use_pose_rope) + ucpe_memory UCPE: PRoPE + absolute-orientation camera encoding +""", + ) + + # Required + parser.add_argument("--ckpt", type=str, required=True, + help="Path to fine-tuned .safetensors checkpoint") + parser.add_argument("--prompt", type=str, required=True, + help="Text prompt describing the scene") + parser.add_argument("--output_path", type=str, required=True, + help="Output video path (.mp4)") + + # Memory selection + parser.add_argument("--memory_type", type=str, default="auto", + choices=ALL_MEMORY_TYPES, + help="Memory type (default: auto-detect from checkpoint path)") + + # Model paths + parser.add_argument("--base_model", type=str, + default=os.environ.get("WAN_BASE_MODEL", ""), + help="Wan2.1 base model directory (default: $WAN_BASE_MODEL)") + parser.add_argument("--tokenizer_path", type=str, default=None, + help="Local tokenizer path (default: /google/umt5-xxl when present)") + + # Context image + parser.add_argument("--context_image", type=str, default=None, + help="Path to first-frame context image (enables context memory)") + parser.add_argument( + "--geometry_memory_video", + type=str, + default=None, + help="TSDF/point-cloud-rendered static condition video for geometry-grounded Spatial Memory", + ) + + # Action control + parser.add_argument("--action_path", type=str, default=None, + help="Path to action JSON file (81-frame camera trajectory)") + + # Generation parameters + parser.add_argument("--height", type=int, default=352) + parser.add_argument("--width", type=int, default=640) + parser.add_argument("--num_frames", type=int, default=81, + help="Number of frames per chunk (default: 81)") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--num_inference_steps", type=int, default=50) + parser.add_argument("--sigma_shift", type=float, default=15.0) + parser.add_argument("--cfg_scale", type=float, default=5.0) + parser.add_argument("--negative_prompt", type=str, default=None, + help="Negative prompt (default: standard quality filter)") + + # Output + parser.add_argument("--fps", type=int, default=15, help="Output video FPS") + + return parser + + +def main(): + args = build_parser().parse_args() + + # ── Validate paths ────────────────────────────────────────────────── + # Heavy imports deferred so --help works without GPU/conda environment + import torch + from PIL import Image + from env.loop_utils import load_pipeline_and_ckpt, DEFAULT_NEGATIVE_PROMPT + from env.run_replay_loop_two_chunk import run_one_chunk, encode_context_frames_per_frame + from diffsynth import save_video + + neg_prompt = args.negative_prompt if args.negative_prompt else DEFAULT_NEGATIVE_PROMPT + + if not args.base_model: + print("ERROR: --base_model or $WAN_BASE_MODEL must be set.", file=sys.stderr) + sys.exit(1) + + dit_path = os.path.join(args.base_model, "diffusion_pytorch_model.safetensors") + text_encoder_path = os.path.join(args.base_model, "models_t5_umt5-xxl-enc-bf16.pth") + vae_path = os.path.join(args.base_model, "Wan2.1_VAE.pth") + tokenizer_path = args.tokenizer_path or os.path.join(args.base_model, "google", "umt5-xxl") + if not os.path.isdir(tokenizer_path): + tokenizer_path = None + + for p in [dit_path, text_encoder_path, vae_path]: + if not os.path.isfile(p): + print(f"ERROR: base model file not found: {p}", file=sys.stderr) + sys.exit(1) + + if not os.path.isfile(args.ckpt): + print(f"ERROR: checkpoint not found: {args.ckpt}", file=sys.stderr) + sys.exit(1) + + # ── Resolve memory profile ────────────────────────────────────────── + profile = resolve_memory_profile(args.memory_type, args.ckpt) + + # ── Load pipeline + checkpoint ────────────────────────────────────── + print(f"[unified_inference] Loading pipeline from {args.base_model}") + print(f"[unified_inference] Loading checkpoint from {args.ckpt}") + pipe = load_pipeline_and_ckpt( + ckpt_path=args.ckpt, + dit_path=dit_path, + text_encoder_path=text_encoder_path, + vae_path=vae_path, + device="cuda", + add_action_attn=False, + action_use_temporal_attention=True, + tokenizer_path=tokenizer_path, + ) + + # ── Apply memory flags ────────────────────────────────────────────── + apply_profile_to_pipe(pipe, profile) + if args.memory_type == "geometry_spatial_mem" or args.geometry_memory_video: + if getattr(pipe, "geometry_spatial_memory_module", None) is None: + print( + "ERROR: geometry Spatial Memory checkpoint does not contain " + "geometry_spatial_memory_module weights.", + file=sys.stderr, + ) + sys.exit(1) + pipe.use_geometry_spatial_memory = True + + # ── Encode context image (if provided) ────────────────────────────── + context_latents = None + context_actions_t = None + num_context_frames = 0 + + if args.context_image: + if not os.path.isfile(args.context_image): + print(f"ERROR: context image not found: {args.context_image}", file=sys.stderr) + sys.exit(1) + + print(f"[unified_inference] Encoding context image: {args.context_image}") + ctx_pil = Image.open(args.context_image).convert("RGB").resize( + (args.width, args.height), Image.LANCZOS + ) + pipe.load_models_to_device(["vae"]) + with torch.no_grad(): + context_latents = encode_context_frames_per_frame( + pipe, [ctx_pil], pipe.device + ) + num_context_frames = 1 + # Identity RT for context frame (no relative pose change) + identity_rt = [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0] + context_actions_t = torch.tensor([identity_rt], dtype=torch.float32) + + geometry_memory_latents = None + if args.geometry_memory_video: + if not os.path.isfile(args.geometry_memory_video): + print( + f"ERROR: geometry memory video not found: {args.geometry_memory_video}", + file=sys.stderr, + ) + sys.exit(1) + import imageio.v3 as iio + + geometry_frames = [ + Image.fromarray(frame).convert("RGB").resize( + (args.width, args.height), + Image.Resampling.LANCZOS, + ) + for frame in iio.imiter(args.geometry_memory_video) + ] + if not geometry_frames: + print("ERROR: geometry memory video contains no frames.", file=sys.stderr) + sys.exit(1) + pipe.load_models_to_device(["vae"]) + with torch.no_grad(): + geometry_video = pipe.preprocess_video(geometry_frames) + if geometry_video.dim() == 4: + geometry_video = geometry_video.unsqueeze(0) + geometry_memory_latents = pipe.vae.encode( + [geometry_video[i] for i in range(geometry_video.shape[0])], + device=pipe.device, + tiled=False, + tile_size=None, + tile_stride=None, + ).to(dtype=pipe.torch_dtype, device=pipe.device) + print( + "[unified_inference] Encoded geometry memory video: " + f"{tuple(geometry_memory_latents.shape)}" + ) + + # ── Generate ──────────────────────────────────────────────────────── + print(f"[unified_inference] Generating {args.num_frames} frames @ {args.width}x{args.height}") + frames = run_one_chunk( + pipe=pipe, + prompt=args.prompt, + use_negative_prompt=neg_prompt, + action_path=args.action_path, + context_latents=context_latents, + num_context_frames=num_context_frames, + context_actions_t=context_actions_t, + geometry_memory_latents=geometry_memory_latents, + chunk_frames=args.num_frames, + h=args.height, + w=args.width, + seed=args.seed, + sigma_shift=args.sigma_shift, + num_inference_steps=args.num_inference_steps, + cfg_scale=args.cfg_scale, + log_prefix="[unified_inference]", + ) + + # ── Save video ────────────────────────────────────────────────────── + os.makedirs(os.path.dirname(os.path.abspath(args.output_path)), exist_ok=True) + save_video(frames, args.output_path, fps=args.fps, quality=5) + print(f"[unified_inference] Video saved to {args.output_path}") + + +if __name__ == "__main__": + main() diff --git a/code/requirements.txt b/code/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..2df1d5f6c1104990ebff3f8421262601f3898524 --- /dev/null +++ b/code/requirements.txt @@ -0,0 +1,24 @@ +accelerate +decord +diffusers +einops +imageio +imageio-ffmpeg +modelscope +numpy<2 +opencv-python +pandas +peft +pillow +protobuf +pyyaml +requests +rotary_embedding_torch +safetensors +scikit-image +scipy +sentencepiece +timm +tqdm +transformers<5 +wandb diff --git a/code/retrieve.md b/code/retrieve.md new file mode 100644 index 0000000000000000000000000000000000000000..96aa041c3ea5cb1550ad343931ae69d8086b3b52 --- /dev/null +++ b/code/retrieve.md @@ -0,0 +1,770 @@ +# Retrieval-Uncertainty Loss & Memory-Retrieval Improvements + +This document gives (1) a concrete, codebase-accurate implementation of the +**retrieval-uncertainty loss** you described, and (2) a set of additional +methods to improve the model's ability to retrieve preceding (context) frames. + +> **Loss you asked for** +> `L_unc = e^(-um) · sg(MSE(VAE(Target View), predicted_x0)) + um` +> where `um` is the model's predicted (log-)uncertainty about the current +> prediction and `sg(·)` is stop-gradient. +> +> This is the **Kendall–Gal heteroscedastic-uncertainty** objective. With the +> MSE stop-gradiented, the uncertainty head learns a *per-token retrieval +> confidence* (low `um` where the model reconstructs the target well, high `um` +> where it fails / "forgets"). That confidence map is then reused to **reweight +> the main denoising loss**, focusing capacity on the tokens the memory pathway +> currently fails to retrieve. + +Read `CLAUDE.md` and `attention.md` first for the two-chunk paradigm and the +public-repo constraints. + +--- + +## 0. Where the loss lives in this codebase + +The training loss is **flow-matching MSE**, computed in +`diffsynth/pipelines/wan_video_new.py → WanVideoPipeline.training_loss(...)`. +The relevant facts (verified against the code): + +- The scheduler is `FlowMatchScheduler` (`diffsynth/schedulers/flow_match.py`): + - `add_noise`: `x_t = (1 - σ)·x0 + σ·noise` + - `training_target`: `v = noise − x0` ← **the model predicts velocity**, not x0 +- In `training_mode == "context"` (the default for memory baselines), the + target tokens are `input_latents` (already VAE-encoded — this **is** + `VAE(Target View)` in latent space), and the loss is: + + ```python + noisy_target_latents = self.scheduler.add_noise(target_latents, target_noise, timestep) # x_t + training_target = self.scheduler.training_target(target_latents, target_noise, timestep) # v = noise - x0 + ... + noise_pred = self.model_fn(**inputs, timestep=timestep) # v_pred over [target | context] tokens + target_noise_pred = noise_pred[:, :, :target_latents.shape[2], :, :] # suffix layout: target first + loss = F.mse_loss(target_noise_pred.float(), training_target.float()) + loss = loss * self.scheduler.training_weight(timestep) + ``` + +**Recovering `predicted_x0`** (needed for your loss) is exact under flow +matching — no extra forward pass: + +``` +x0_pred = x_t − σ · v_pred = noisy_target_latents − σ · target_noise_pred +``` + +So `MSE(VAE(Target View), predicted_x0) = MSE(target_latents, x0_pred)`, computed +entirely in latent space — no VAE decode is needed during training. + +> `σ` for the sampled `timestep` is `self.scheduler.sigmas[timestep_id]` where +> `timestep_id = argmin(|scheduler.timesteps − timestep|)`. There are two +> equivalent target-token layouts (`context_position == "suffix"` vs `prefix`); +> the slice that selects target tokens is already computed as +> `target_noise_pred` — reuse it. + +--- + +## 1. The module (already added) + +`diffsynth/models/memory/uncertainty.py` (exported from +`diffsynth/models/memory/__init__.py`). Key pieces: + +- `UncertaintyHead(in_channels)` — a 1×1×1 Conv3d MLP mapping the per-token x0 + prediction `(B, C, T, H, W)` → log-uncertainty `um` `(B, 1, T, H, W)`. The + final conv is **zero-initialised**, so `um ≡ 0` at init ⇒ `exp(−um) ≡ 1` + (a no-op weighting) ⇒ the existing training dynamics are unchanged on step 0. +- `recover_x0_flow_match(noisy_target, v_pred, sigma)` — `x_t − σ·v`. +- `per_token_mse_map(x0_pred, x0_target)` — channel-mean MSE → `(B,1,T,H,W)`. +- `retrieval_uncertainty_loss(um, mse_map, detach_mse=True)` — returns + `mean(exp(−um)·sg(MSE) + um)`. + +Smoke-tested: `um` is `(B,1,T,H,W)`, `==0` at init, loss is finite, gradients +flow into the head. + +--- + +## 2. Wiring it into training + +Five small edits. All are additive and gated by a new flag so default behaviour +is untouched. + +### 2a. `diffsynth/pipelines/wan_video_new.py` — `WanVideoPipeline.__init__` + +Create the head lazily (the latent channel count `C` is known from the DiT +`in_dim`, typically 16 for Wan 2.1). Add near where other memory attributes are +set: + +```python +self.use_retrieval_uncertainty = False +self.retrieval_uncertainty_weight = 1.0 +self.uncertainty_head = None # nn.Module, built on first use +``` + +### 2b. `WanVideoPipeline.training_loss` — compute the extra loss + +In **both** the `"context"` and `"predict"` branches, right after the existing +`loss = ... * training_weight(timestep)`, insert: + +```python +if getattr(self, "use_retrieval_uncertainty", False): + from diffsynth.models.memory.uncertainty import ( + UncertaintyHead, recover_x0_flow_match, per_token_mse_map, + retrieval_uncertainty_loss, + ) + # σ for this timestep (flow-match scheduler). + sched = self.scheduler + timestep_id = torch.argmin( + (sched.timesteps - timestep.to(sched.timesteps.device)).abs()) + sigma = sched.sigmas[timestep_id].to(target_noise_pred.dtype) + + # x_t and v_pred for the TARGET tokens only (reuse the existing slice). + x_t_target = noisy_target_latents # (B,C,T,H,W) + x0_pred = recover_x0_flow_match(x_t_target, target_noise_pred, sigma) + x0_target = target_latents # = VAE(Target View) + + # Lazily build the head with the right channel count, on the right device. + if self.uncertainty_head is None: + self.uncertainty_head = UncertaintyHead(x0_pred.shape[1]).to( + device=x0_pred.device, dtype=torch.float32) + um = self.uncertainty_head(x0_pred) # (B,1,T,H,W) + + mse_map = per_token_mse_map(x0_pred, x0_target) + loss_unc = retrieval_uncertainty_loss(um, mse_map, detach_mse=True) + + # (Optional, Method A) reweight the MAIN denoising loss by confidence so the + # denoiser/memory pathway focuses on hard-to-retrieve tokens. Detach um here + # so this term trains the denoiser, not the head. + # per_token_main = (target_noise_pred.float() - training_target.float()).pow(2).mean(1, keepdim=True) + # w = torch.exp(-um.detach()).clamp(0.1, 10.0) + # loss = (w * per_token_main).mean() * self.scheduler.training_weight(timestep) + + loss = loss + self.retrieval_uncertainty_weight * loss_unc +``` + +> `target_latents`, `noisy_target_latents`, `target_noise_pred`, and `timestep` +> all already exist as locals in that scope — no signature changes needed. + +### 2c. `src/model_training/train.py` — argparse flags + +Add to the numeric-default tuples (near `--timestep_shift`): + +```python +("--retrieval_uncertainty_weight", dict(type=float, default=1.0)), +``` + +and add `"--use_retrieval_uncertainty"` to the list of store-true flags +(alongside `"--use_block_wise_ssm"`, ~line 1510). + +### 2d. `src/model_training/train.py` — push flags onto the pipe + +In the trainer `__init__` (where `self.pipe.use_spatial_memory = ...` is set, +~line 984), add: + +```python +self.pipe.use_retrieval_uncertainty = bool(use_retrieval_uncertainty) +self.pipe.retrieval_uncertainty_weight = float(retrieval_uncertainty_weight) +``` + +and thread the two values in from `_arg(...)` at the trainer construction call +(~line 1671), mirroring `timestep_shift`. + +### 2e. Make the head trainable **and saved** + +The optimizer collects `model.trainable_modules()` = all params with +`requires_grad=True`, and `--save_full_model` exports the whole DiT state +(otherwise only `requires_grad` params are exported via +`export_trainable_state_dict`). The head lives on `self.pipe`, **not** inside +`dit`, so: + +- Its params are created with `requires_grad=True` by default ✔ (so AdamW will + pick them up **provided the optimizer is built after the head exists**). The + head is built lazily on the first `training_loss` call, which is *after* + `torch.optim.AdamW(model.trainable_modules(), ...)` (~line 1794). **Fix:** + build the head eagerly so it is registered before the optimizer is created — + add this right after the block-replacement section in `train.py`: + + ```python + if _arg('use_retrieval_uncertainty', False): + from diffsynth.models.memory.uncertainty import UncertaintyHead + _c = int(getattr(model.pipe.dit, "in_dim", 16)) + model.pipe.uncertainty_head = UncertaintyHead(_c).to( + device=next(model.pipe.dit.parameters()).device, dtype=torch.float32) + ``` + +- For checkpointing, the head is an attribute of `self.pipe`, which is a + submodule of the training `model`, so `accelerator.get_state_dict(model)` + includes `pipe.uncertainty_head.*` keys. With `--save_full_model` they are + saved; the keys are ignored at inference (the head is not needed to generate). + If you do **not** use `--save_full_model`, ensure the head params have + `requires_grad=True` (they do) so `export_trainable_state_dict` keeps them. + +### 2f. Launcher + +Copy an existing memory launcher (e.g. +`train/memory_baselines_basic/run_spatial_memory_baseline.sh`) and add: + +```bash + --use_retrieval_uncertainty --retrieval_uncertainty_weight 1.0 \ +``` + +Keep every other hyperparameter identical to the baseline row you are comparing +against — this is a controlled ablation; only the loss should change. + +### 2g. Sanity check before a full run + +```bash +PYTHONPATH=. python3 tests/test_two_chunk_anchor_readout.py +# plus a tiny head test (identity-at-init + grad flow), e.g.: +PYTHONPATH=. python3 - <<'PY' +import importlib.util, torch +s=importlib.util.spec_from_file_location('u','diffsynth/models/memory/uncertainty.py') +u=importlib.util.module_from_spec(s); s.loader.exec_module(u) +h=u.UncertaintyHead(16); xt=torch.randn(1,16,21,44,80); v=torch.randn_like(xt) +x0=u.recover_x0_flow_match(xt,v,torch.tensor(.7)); um=h(x0) +assert float(um.abs().max())==0.0 # identity at init +L=u.retrieval_uncertainty_loss(um,u.per_token_mse_map(x0,torch.randn_like(x0))) +L.backward(); assert h.net[0].weight.grad is not None +print("ok", float(L)) +PY +``` + +--- + +## 3. Why this helps retrieval (and how to read the signal) + +`um` becomes a learned, per-token map of **where the model fails to reconstruct +the target from memory**. Two ways to exploit it: + +- **Diagnostic** — log `um` heatmaps to W&B alongside the two-chunk + left/right-rotation monitor (`--sampling_atomic_left_right`). High-`um` + regions on the revisit tail localise *what* the memory is dropping (object + identity vs background vs camera geometry). +- **Loss reweighting (Method A above)** — `exp(−um.detach())` upweights the + main denoising loss on tokens the model is *confident and wrong* about, + pushing the memory pathway to fix systematic retrieval failures rather than + averaging error uniformly. + +--- + +## 4. Additional methods to improve preceding-frame retrieval + +Ordered roughly by expected impact / effort. All are compatible with the +two-chunk setup and the existing memory families. + +### Method A — Confidence-reweighted denoising loss +Already sketched in §2b. Uses `exp(−um.detach())` to focus the **denoiser** on +hard-to-retrieve tokens. Cheap, synergises directly with the uncertainty head. + +### Method B — Explicit retrieval-consistency (anchor) loss +Add a term that directly penalises drift between the **first/anchor frame** and +the **revisit tail** in latent space, since revisit consistency is exactly what +the paper measures. After recovering `x0_pred` for the target tokens: + +``` +L_anchor = MSE( x0_pred[revisit_tail_tokens], context_latents[anchor_token] ) +``` +restricted to samples where the trajectory returns near the start pose (the +codebase already constructs loop-closure probes; reuse +`env/loop_utils.py` / the `replay` context source to identify revisit tokens). +This trains the memory pathway to *reproduce* stored content, not merely to +denoise plausibly. Gate behind `--use_anchor_consistency_loss`. + +### Method C — Contrastive memory read-out (InfoNCE) +Make the memory read-out **discriminative**: the target token's retrieved +memory feature should match its *own* context frame more than other frames'. +Take per-frame pooled features from the context tokens (before they enter the +DiT blocks) and the corresponding target query features, and add an InfoNCE +loss pulling matched (target-frame ↔ source-frame) pairs together and pushing +mismatched pairs apart. This sharpens *which* preceding frame is retrieved — +particularly useful for the Spatial and Context-K families. Implement as a +small head reading the block hidden state (same hook point as block-wise SSM in +`DiTBlock_w_Action`, see `attention.md`). + +### Method D — Harder/longer context sampling (curriculum) +Retrieval is only as good as the supervision distribution. Levers already in +the data path: +- Increase `--context_memory_frames` (K) and/or widen the temporal gap between + context and target so the model must retrieve *distant* history, not adjacent + frames (`--context_source replay`, `--prev_chunk_frames`). +- Curriculum: start with short gaps, anneal to longer gaps over training. +- Mix revisit-style samples (leave-and-return) more heavily — the two-chunk + `--sampling_atomic_left_right` probe shows what to oversample. +Pure data/schedule change; no model edits. + +### Method E — Memory dropout / robustness regularisation +Randomly drop or noise a subset of context tokens during training +(`--context_drop_prob`, `--context_noise_std` already exist). Forcing the model +to retrieve from partial memory improves robustness and prevents trivial +copy-through, which tends to help long-horizon revisit. Tune these existing +flags rather than adding code. + +### Method F — Cross-attention readout supervision for Spatial memory +For the spatial family (`spatial_cross_attn_readout`), add an auxiliary loss +that encourages the read-out attention map to concentrate on the spatially +corresponding stored region (when camera RT gives a known correspondence). +This is a targeted version of Method C for the spatial grid memory. + +### Recommended first experiment +1. Implement §1–§2 (uncertainty head + loss), train one row vs its baseline. +2. Turn on **Method A** (confidence reweighting) — likely the largest gain per + line of code. +3. Add **Method B** (anchor consistency) if revisit MSE is still the bottleneck. +Evaluate all with the existing tiers: +```bash +export CKPT=outputs//epoch-0.safetensors +bash eval/v2/run_basic_replay_gt.sh +bash eval/v2/run_static_consistency_loop_and_revisit.sh +PHASE=stage1 OOD_DIR=assets/opendomain_revisit bash eval/v2/revisit_suite/run_one_click_revisit_eval.sh +``` +Compare revisit-tail MSE / PSNR / LPIPS against the unmodified baseline row. + +--- + +## 5. Pitfalls + +- **Predicting x0 vs velocity.** The model outputs **velocity** `v = noise − x0`. + Do **not** feed `target_noise_pred` directly as `x0` — always recover via + `x0 = x_t − σ·v` (`recover_x0_flow_match`). Getting this wrong silently + inverts the uncertainty signal. +- **Non-zero `um` at init.** Keep the head's final layer zero-initialised; a + non-zero init multiplies the main loss by an arbitrary factor on step 0 and + destabilises early training. +- **Optimizer misses the head.** Build the head **before** + `AdamW(model.trainable_modules())` is constructed (see §2e), or its params + won't be optimised. +- **Stop-gradient.** With `detach_mse=True` the uncertainty term trains only the + head. If you want it to also shape the denoiser, use Method A's + `exp(−um.detach())` reweighting of the main loss — don't simply drop the + stop-gradient on the MSE (that lets the model lower the loss by inflating + `um`, i.e. "predict badly on purpose"). +- **dtype/autocast.** Compute the head and the loss in fp32 (the module already + casts), and the `um` clamp keeps `exp(−um)` finite under bf16 autocast. +- **Public-repo constraints** (`CLAUDE.md`): no machine-local paths, minimal + diffs, don't commit `outputs/` or weights. + +--- + +## 6. Where the retrieval target comes from (what "correct retrieval" means) + +A confidence map is only meaningful relative to a **target** that defines correct +retrieval. There is no single target — there is a hierarchy of increasingly +strict definitions, and which one you pick decides what your confidence map +actually measures. This codebase already computes the geometric ones. + +### Background: what is RT? + +**RT = Rotation + Translation = the camera extrinsics** (the rigid-body pose of +the camera). In this repo an RT is a **12-dim row-major vector** +`[t_x, t_y, t_z, R_11, R_12, R_13, R_21, R_22, R_23, R_31, R_32, R_33]` +— a 3×1 translation `t` followed by a flattened 3×3 rotation `R` +(`src/model_training/rt_utils.py` docstring). The `MLP_CamPose(pose_dim=12)` +inside `DiTBlock_w_Action` consumes exactly this 12-vector per latent frame. + +The **relative RT** between a context frame *i* and the reference (target) frame +maps points from one camera frame into the other — this is what lets you +*reproject* context content into the target view. It is computed by +`rt_utils.convert_rt_to_relative(rt_list_all, ref_rt)`: + +``` +R_rel = R_ref⁻¹ · R_i , t_rel = R_ref⁻¹ · t_i + (−R_ref⁻¹ · t_ref) +``` + +(`R_ref⁻¹ = R_refᵀ` since rotations are orthonormal). Camera poses come from the +per-frame JSONs via `pose_to_rt(pose)` (paper default: XY translation + Z-axis +yaw only). Enabled in training by `--use_rt_relative` (env `USE_RT_RELATIVE`). + +### Level 0 — Reconstruction target (what the base loss already uses) + +Weakest definition: *correct retrieval = the token was denoised well.* Target is +`x0_target = target_latents = VAE(Target View)`; supervision is the per-token MSE +already computed in `training_loss`. **Limitation:** it conflates two failure +modes — (a) the memory pathway *failed to retrieve* the right context, vs. +(b) the content is *genuinely novel / newly revealed* and no memory could help. +For a *retrieval* confidence map you want to isolate (a), so Level 0 alone is the +wrong target. + +### Level 1 — Geometric co-visibility (already in the repo) + +The retrieval target here is **not learned** — it is a precomputed geometric +label answering *"which past frames actually share field-of-view with the +current frame?"*: + +- **`overlap_labels/{video_name}/{frame_idx}.json`** → + `{"overlapping_frames": ["2796", "2797", ...]}`. For each frame, the list of + historical frames that co-observe the same scene. **This is the ground-truth + retrieval target at frame granularity.** Loaded by + `fov_retrieval.load_overlap_frames()` and consumed by + `fov_training_integration.retrieve_fov_context_frames()` to *select* the + context frames during training. +- **`fov_retrieval.compute_fov_overlap_3d(pose1, pose2, fov=52.67°)`** computes a + continuous overlap score in [0,1] from the 6-DoF poses (mutual visibility + + forward-direction similarity). This is the function that *generates* the + labels. + +Use it as a **target for the confidence head**: a token whose co-visible content +the model reproduced ⇒ confidence high; a token that *had* co-visible support +but was reproduced wrong ⇒ retrieval failure (what you want to flag). The +co-visibility set also gives a **mask**: only ask "did you retrieve correctly?" +where retrieval was geometrically possible. + +### Level 2 — Reprojection correspondence (the per-token target you asked for) + +Level 1 is frame-level + coarse-region. Level 2 tightens it to **per latent +token**: use the relative RT to warp the co-visible context latent into the +current view, then define correct retrieval token-by-token. This removes the +Level-0 ambiguity (novel regions are masked out of the retrieval loss). + +#### 6.1 Token ↔ pixel geometry in this stack (must get this right) + +To reproject *into latent-token space* you need the compression factors: + +| Stage | Factor | Source | +| --- | --- | --- | +| VAE spatial downsample | **÷8** (three 2× `downsample2d/3d` blocks) | `wan_video_vae.py` Resample blocks | +| VAE temporal downsample | **÷4** (`temperal_downsample=[True,True,False]`, +1 for the first frame) | `wan_video_vae.py:284` | +| DiT patchify | **(1, 2, 2)** | `wan_video_dit.py:511` `patch_size=(1,2,2)` | + +So one latent **token** covers a `8·2 = 16` px × `16` px region of the original +frame (spatially), and the latent grid for a `352×640` frame is +`H_lat = 352/8 = 44`, `W_lat = 640/8 = 80`, then patchified by 2 → +`22 × 40` token grid per latent frame. **Reproject at the latent-pixel grid +(44×80), then patch-pool to the token grid (22×40)** to match `um`'s resolution. + +#### 6.2 Building the per-token reprojection-confidence map + +Given a context frame *i* and the target frame, with relative pose `(R_rel, +t_rel)` from `convert_rt_to_relative`, the homltography/flow that maps target +latent-pixel `(u,v)` ↔ context latent-pixel depends on scene depth. Two regimes: + +- **Depth available** (SpatialVID has more geometry than the static pool): + full reprojection `p_ctx = K · (R_rel · (depth · K⁻¹ · p_tgt) + t_rel)`. +- **No depth / planar approximation** (static pool, paper's XY+yaw setting): a + **homography** `Hℓ` suffices because motion is dominated by yaw + translation + on a plane — exactly the regime `pose_to_rt(constrain_to_xy=True)` encodes. + Build `Hℓ` from `(R_rel, t_rel)` and a reference plane normal/depth. + +The confidence target is then: *warp the context latent into the target view and +measure how close the model's `x0_pred` is to that warped evidence, only where +co-visibility holds.* + +```python +# diffsynth/models/memory/reproj_confidence.py (sketch — add as a new module) +import torch +import torch.nn.functional as F + + +def latent_grid_hw(height_px: int, width_px: int): + """Latent-pixel grid before DiT patchify: VAE divides spatial by 8.""" + return height_px // 8, width_px // 8 + + +def warp_context_latent(ctx_latent, H_rel): + """ + Warp a context latent frame into the target view via a 3x3 homography H_rel + expressed in *latent-pixel* coordinates (44x80 for 352x640). + + ctx_latent: (B, C, Hl, Wl) single context latent frame + H_rel: (B, 3, 3) target-latent-pixel -> context-latent-pixel + returns: (B, C, Hl, Wl) warped context, (B,1,Hl,Wl) valid mask + """ + B, C, Hl, Wl = ctx_latent.shape + ys, xs = torch.meshgrid( + torch.arange(Hl, device=ctx_latent.device, dtype=torch.float32), + torch.arange(Wl, device=ctx_latent.device, dtype=torch.float32), + indexing="ij", + ) + ones = torch.ones_like(xs) + grid = torch.stack([xs, ys, ones], dim=-1).reshape(1, Hl * Wl, 3).expand(B, -1, -1) + + src = torch.bmm(grid, H_rel.transpose(1, 2)) # (B, Hl*Wl, 3) + src = src[..., :2] / src[..., 2:3].clamp(min=1e-6) # homogeneous divide + sx, sy = src[..., 0], src[..., 1] + + # Normalise to grid_sample's [-1, 1] coordinates. + gx = (sx / (Wl - 1)) * 2 - 1 + gy = (sy / (Hl - 1)) * 2 - 1 + samp = torch.stack([gx, gy], dim=-1).reshape(B, Hl, Wl, 2) + + warped = F.grid_sample(ctx_latent, samp, mode="bilinear", + padding_mode="zeros", align_corners=True) + valid = ((gx >= -1) & (gx <= 1) & (gy >= -1) & (gy <= 1)).float() + return warped, valid.reshape(B, 1, Hl, Wl) # in-FOV co-visibility mask + + +def reprojection_confidence_map(x0_pred, ctx_latents, H_rels, + patch=2, tau=1.0): + """ + Per-token retrieval-confidence target from reprojection correspondence. + + x0_pred: (B, C, T, Hl, Wl) recovered x0 for TARGET tokens (per latent frame) + ctx_latents: (B, C, K, Hl, Wl) clean context latents (VAE-encoded history) + H_rels: (B, T, K, 3, 3) target-frame t <- context-frame k homographies + (latent-pixel coords), from convert_rt_to_relative + Returns: + conf_tok: (B, 1, T, Hl//patch, Wl//patch) in [0,1], token-resolution + mask_tok: (B, 1, T, Hl//patch, Wl//patch) co-visibility (any context covers token) + """ + B, C, T, Hl, Wl = x0_pred.shape + K = ctx_latents.shape[2] + best_err = x0_pred.new_full((B, 1, T, Hl, Wl), float("inf")) + any_valid = x0_pred.new_zeros((B, 1, T, Hl, Wl)) + + for t in range(T): + for k in range(K): + warped, valid = warp_context_latent(ctx_latents[:, :, k], H_rels[:, t, k]) + err = (x0_pred[:, :, t] - warped).pow(2).mean(dim=1, keepdim=True) # (B,1,Hl,Wl) + err = torch.where(valid > 0, err, best_err[:, :, t]) + best_err[:, :, t] = torch.minimum(best_err[:, :, t], err) # best matching ctx frame + any_valid[:, :, t] = torch.maximum(any_valid[:, :, t], valid) + + best_err = torch.where(torch.isfinite(best_err), best_err, torch.zeros_like(best_err)) + conf = torch.exp(-best_err / tau) # low reprojection error -> high confidence + conf = conf * any_valid # undefined where nothing is co-visible + + # Patch-pool latent-pixel grid (44x80) down to token grid (22x40) to match `um`. + conf_tok = F.avg_pool3d(conf, kernel_size=(1, patch, patch)) + mask_tok = (F.avg_pool3d(any_valid, kernel_size=(1, patch, patch)) > 0).float() + return conf_tok, mask_tok +``` + +**Where `H_rels` comes from.** In `training_loss` you already have (or can pass +through `inputs`) the per-frame RTs. For target latent frame *t* and context +frame *k*: `rel = convert_rt_to_relative([rt_k], ref_rt=rt_t)[0]`, parse into +`(R_rel, t_rel)`, and convert to a latent-pixel homography with the intrinsics +scaled by 1/8 (latent) — under the paper's XY+yaw planar setting a homography is +the correct first-order model. Precompute `H_rels` on CPU/numpy in the dataloader +(the RTs are already loaded for the action MLP) and pass them in as a tensor; +avoid per-step Python geometry in the hot loop. + +#### 6.3 Two ways to use the reprojection map + +This is the bridge to Category B / Category C from the discussion: + +- **(B) As a near-ground-truth confidence map directly** — `conf_tok` *is* a + retrieval-confidence map, no learning required. Use it to reweight the main + loss (`w = conf_tok` upweights tokens that *should* be retrievable, focusing + the memory pathway on co-visible content) or as an eval-time diagnostic over + the revisit tail. + +- **(C) As the supervision target for a predictive head** — train the + `UncertaintyHead` (or a dedicated `ConfidenceHead`) to **predict `conf_tok` + before generation**, supervised only on `mask_tok` tokens: + + ```python + pred_conf = torch.sigmoid(-um) # head's confidence in [0,1] + loss_conf = (mask_tok * (pred_conf - conf_tok.detach()).pow(2)).sum() \ + / mask_tok.sum().clamp(min=1.0) + loss = loss + lambda_conf * loss_conf + ``` + + This gives a **calibrated, forward-time** confidence signal grounded in + geometry, instead of the self-supervised heteroscedastic target — and it + cleanly answers "do we know, during training, whether we retrieved the correct + context?": *yes, because geometry tells us which tokens had retrievable + support and reprojection tells us whether the model reproduced it.* + +### How to know, during training, if retrieval is correct — summary + +| Target level | Source in repo | "Correct" means | Strength / caveat | +| --- | --- | --- | --- | +| **0 Reconstruction** | `target_latents` (base loss) | low token MSE | weak — confounds forgetting vs. novelty | +| **1 Co-visibility** | `overlap_labels/*.json`, `compute_fov_overlap_3d` | model uses the geometrically co-visible frames | frame/region-level; FOV-frustum, **not depth/occlusion aware** | +| **2 Reprojection** | relative RT (`convert_rt_to_relative`) + warp | warped co-visible evidence matches `x0_pred`, masked to co-visible tokens | per-token, strongest; needs depth or planar/homography assumption | + +**Honesty caveat (state this in any writeup):** the overlap labels and +`compute_fov_overlap_3d` are **camera-frustum** co-visibility from poses, *not* +depth-aware occlusion — two frames can be marked co-visible when an occluder +blocks the shared content. Level-2 reprojection inherits this: the homography +regime assumes near-planar / yaw-dominant motion (the paper's `constrain_to_xy` +setting). For truly metric per-token correspondence, use depth (better available +in the SpatialVID dynamic pool) and a full reprojection rather than a homography. + +--- + +## 7. Depth-aware per-token confidence (the accurate version) + +With depth you replace the §6.2 **homography** (planar, yaw-dominant +approximation) by a **full metric reprojection with occlusion reasoning**. This +removes the two failure modes of the homography path: (i) it handles arbitrary +3D scene geometry and 6-DoF motion, not just a reference plane, and (ii) it can +*detect occlusion* — telling apart "co-visible and the model retrieved it" from +"the frustum overlaps but an occluder hides the content" (the exact blind spot +of the FOV-frustum labels in §6, Level 1). + +### 7.1 Pose convention in this repo (get the direction right) + +`fov_retrieval.compute_fov_overlap_3d` treats `position` as the **camera centre +in world coordinates** `C` and the third column `R[:,2]` as the world-space +**forward** axis. So the stored 12-dim RT `[t | R]` is **camera-to-world**: + +``` +X_world = R · X_cam + C # R = R_cam→world, C = camera centre = t +X_cam = Rᵀ · (X_world − C) # world → camera (inverse) +``` + +(`R⁻¹ = Rᵀ` for a rotation). This is the opposite direction from a +"world-to-camera extrinsic" `[R|t]` convention — using the wrong one silently +flips the reprojection, so anchor on this. + +### 7.2 The reprojection (target token → 3D → context frame) + +For a target latent token at pixel `p_t=(u,v)` in latent-pixel coords with +metric depth `d`: + +1. **Back-project to the target camera ray, scale by depth:** + `X_cam_t = d · K⁻¹ · [u, v, 1]ᵀ` +2. **Target camera → world** (camera-to-world): + `X_world = R_t · X_cam_t + C_t` +3. **World → context camera k:** + `X_cam_k = R_kᵀ · (X_world − C_k)` +4. **Project into context frame k:** + `p_k = K · X_cam_k / z_k`, where `z_k = X_cam_k.z` + +`K` is the **latent-resolution** intrinsic: build it from the FOV +(`fov=52.67°`, the same constant `compute_fov_overlap_3d` uses) and divide focal +length + principal point by the VAE spatial factor 8 (so it acts on the 44×80 +latent-pixel grid, matching §6.1). Then sample the context latent at `p_k` and, +crucially, **also sample the context depth at `p_k`** for the occlusion test. + +### 7.3 Occlusion test (what depth buys you) + +A target point is genuinely visible in context frame k only if its reprojected +depth `z_k` matches the context frame's own recorded depth at `p_k`. If the +context depth is *closer* than `z_k`, something else occludes the point — mark +it **not co-visible** even though the frustum overlaps: + +``` +visible_k = (z_k ≤ depth_ctx_k(p_k) · (1 + occ_thresh)) +``` + +This is a forward z-buffer check (`occ_thresh` ~0.05–0.1 absorbs depth noise). +It is exactly the discriminator the §6 Level-1 labels lack. + +### 7.4 Sketch (syntax-checked) + +```python +# diffsynth/models/memory/reproj_confidence_depth.py (sketch) +import torch +import torch.nn.functional as F + + +def intrinsics_latent(width_px, height_px, fov_deg=52.67, vae_down=8): + """Latent-resolution pinhole intrinsics K (focal & principal point ÷ VAE factor).""" + import math + Wl, Hl = width_px // vae_down, height_px // vae_down + f_px = (width_px / 2.0) / math.tan(math.radians(fov_deg) / 2.0) + f_lat = f_px / vae_down + K = torch.tensor([[f_lat, 0.0, Wl / 2.0], + [0.0, f_lat, Hl / 2.0], + [0.0, 0.0, 1.0]]) + return K, Hl, Wl + + +def reproject_target_to_context(depth_t, R_t, C_t, R_k, C_k, K, Kinv): + """ + Map every target latent-pixel into context frame k via depth + camera-to-world RT. + depth_t: (B,1,Hl,Wl) metric depth of TARGET latent frame + R_t,R_k: (B,3,3) camera->world rotations ; C_t,C_k: (B,3) camera centres + returns: grid (B,Hl,Wl,2) for grid_sample, in_fov mask (B,1,Hl,Wl), + z_k_map (B,1,Hl,Wl) reprojected depth in context camera + """ + B, _, Hl, Wl = depth_t.shape + dev = depth_t.device + ys, xs = torch.meshgrid(torch.arange(Hl, device=dev, dtype=torch.float32), + torch.arange(Wl, device=dev, dtype=torch.float32), + indexing="ij") + ones = torch.ones_like(xs) + pix = torch.stack([xs, ys, ones], -1).reshape(1, Hl * Wl, 3).expand(B, -1, -1) + ray = torch.bmm(pix, Kinv.transpose(1, 2)) # K^-1 [u,v,1] + d = depth_t.reshape(B, Hl * Wl, 1) + Xc_t = ray * d # target camera coords + Xw = torch.bmm(Xc_t, R_t.transpose(1, 2)) + C_t.reshape(B, 1, 3) # cam->world + Xc_k = torch.bmm(Xw - C_k.reshape(B, 1, 3), R_k) # world->context cam (R_k^T via right-mul) + z_k = Xc_k[..., 2:3].clamp(min=1e-6) + proj = torch.bmm(Xc_k / z_k, K.transpose(1, 2)) + u, v = proj[..., 0], proj[..., 1] + gx = (u / (Wl - 1)) * 2 - 1 + gy = (v / (Hl - 1)) * 2 - 1 + grid = torch.stack([gx, gy], -1).reshape(B, Hl, Wl, 2) + in_fov = ((gx >= -1) & (gx <= 1) & (gy >= -1) & (gy <= 1)).float().reshape(B, 1, Hl, Wl) + return grid, in_fov, z_k.reshape(B, 1, Hl, Wl) + + +def depth_aware_confidence(x0_pred, ctx_latents, ctx_depths, depth_t, + R_t, C_t, R_k_list, C_k_list, K, + patch=2, tau=1.0, occ_thresh=0.1): + """ + Per-token retrieval-confidence map via metric reprojection + occlusion test. + x0_pred: (B,C,T,Hl,Wl) recovered x0 for target tokens + ctx_latents:(B,C,K,Hl,Wl) clean context latents ; ctx_depths:(B,1,K,Hl,Wl) + depth_t: (B,1,T,Hl,Wl) target-frame metric depth + R_t,C_t: (B,T,3,3),(B,T,3) target cam->world per latent frame + R_k_list,C_k_list: lists of (B,3,3),(B,3) per context frame + """ + B, C, T, Hl, Wl = x0_pred.shape + Kb = K.unsqueeze(0).expand(B, -1, -1) + Kinv = torch.inverse(K).unsqueeze(0).expand(B, -1, -1) + best_err = x0_pred.new_full((B, 1, T, Hl, Wl), float("inf")) + any_valid = x0_pred.new_zeros((B, 1, T, Hl, Wl)) + for t in range(T): + for k in range(len(R_k_list)): + grid, in_fov, z_proj = reproject_target_to_context( + depth_t[:, :, t], R_t[:, t], C_t[:, t], R_k_list[k], C_k_list[k], Kb, Kinv) + warped = F.grid_sample(ctx_latents[:, :, k], grid, mode="bilinear", + padding_mode="zeros", align_corners=True) + ctx_z = F.grid_sample(ctx_depths[:, :, k], grid, mode="bilinear", + padding_mode="zeros", align_corners=True) + visible = (z_proj <= ctx_z * (1.0 + occ_thresh)).float() # z-buffer occlusion test + valid = in_fov * visible + err = (x0_pred[:, :, t] - warped).pow(2).mean(1, keepdim=True) + err = torch.where(valid > 0, err, best_err[:, :, t]) + best_err[:, :, t] = torch.minimum(best_err[:, :, t], err) + any_valid[:, :, t] = torch.maximum(any_valid[:, :, t], valid) + best_err = torch.where(torch.isfinite(best_err), best_err, torch.zeros_like(best_err)) + conf = torch.exp(-best_err / tau) * any_valid + conf_tok = F.avg_pool3d(conf, (1, patch, patch)) + mask_tok = (F.avg_pool3d(any_valid, (1, patch, patch)) > 0).float() + return conf_tok, mask_tok +``` + +### 7.5 Getting depth into latent-token space + +- **Source.** The SpatialVID dynamic pool carries richer geometry than the + static pool; if per-frame metric depth is not already exported, run a monocular + depth estimator offline and cache it (do **not** add it to the training hot + loop). The static in-domain pool only has camera poses, so depth-aware + confidence is primarily a **dynamic-pool** technique. +- **Resolution.** Downsample depth to the **latent-pixel grid** (÷8 → 44×80) by + *area/min pooling* (min-pool preserves near surfaces for the occlusion test; + avoid bilinear across depth discontinuities, which invents mid-air depths). +- **Scale.** Metric consistency matters — `z_k` (reprojected) and + `depth_ctx_k` must be in the **same units**. If depth is up-to-scale + (monocular), fit a per-video scale so it is consistent with the RT translation + units, or make the occlusion test **relative** (compare normalised depth + ranks) instead of absolute. +- **Plumbing.** Precompute and pass `depth_t`, `ctx_depths`, and the per-frame + `(R, C)` through `inputs` (the RTs are already loaded for the action MLP — see + §6). Keep the double loop over `T×K` out of the innermost step by vectorising + over `k`, or restrict `k` to the top-N co-visible frames from the §6 Level-1 + overlap labels (cheaper and removes obviously-irrelevant frames first). + +### 7.6 Accuracy ladder (how the targets compare) + +| Variant | Geometry model | Occlusion | Needs | Accuracy | +| --- | --- | --- | --- | --- | +| §6 Level-1 co-visibility | camera frustum (poses only) | ✗ | poses | frame/region | +| §6.2 homography | planar / yaw-dominant | ✗ | poses + plane | per-token, approx | +| **§7 depth reprojection** | full 6-DoF metric | **✓ (z-buffer)** | poses + **depth** | **per-token, metric** | + +The depth-aware map plugs into the **same two consumers** as §6.3: use `conf_tok` +directly to reweight the main loss, or as the supervision target for a +forward-time `ConfidenceHead` (masked on `mask_tok`). The only change is a +strictly more accurate, occlusion-aware target. + +**Caveats specific to depth.** Reprojection confidence is now bounded by *depth +quality*: noisy/biased monocular depth produces false occlusions and warp +errors. Mitigate with a tolerant `occ_thresh`, min-pooled latent depth, and — +when in doubt — fall back to the §6.2 homography or §6 Level-1 mask for frames +whose depth is low-confidence. Dynamic/independently-moving objects also break +the static-scene assumption of any reprojection (the point moved between +frames); mask known-dynamic regions out of the retrieval loss where you can +detect them. + diff --git a/code/scripts/README.md b/code/scripts/README.md new file mode 100644 index 0000000000000000000000000000000000000000..22d8d7076cb0e7162f339e399f1e28d4241c6142 --- /dev/null +++ b/code/scripts/README.md @@ -0,0 +1,55 @@ +# Data scripts (static in-domain pool) + +Scripts for the **static in-domain pool**. Complete download and layout verification first — see **[doc/dataset_preprocessing.md](../doc/dataset_preprocessing.md)**. + +Both training pools share the Echo-Memory layout (`frames/`, `jsons/`, `overlap_labels/`, `metadata_full.csv`). Set `DATASET_BASE_PATH` to the pool root before running these scripts. + +## Metadata CSV + +The Echo-Team packaged static pool already includes `metadata_full.csv`. If you downloaded the upstream pool without metadata, fetch the released metadata: + +```bash +export DATASET_BASE_PATH=data/Context-as-Memory-Dataset +huggingface-cli download Echo-Team/Echo-Memory-Data metadata_full.csv \ + --repo-type dataset \ + --local-dir "${DATASET_BASE_PATH}" +``` + +Or regenerate it locally after changing the pool: + +```bash +bash scripts/run_generate_metadata.sh +``` + +For a smaller custom metadata file, set an output path and row limit: + +```bash +OUTPUT_CSV="${DATASET_BASE_PATH}/metadata_1000.csv" \ +METADATA_MAX_ROWS=1000 \ +bash scripts/run_generate_metadata.sh +``` + +Use that file with `--dataset_metadata_path "${DATASET_BASE_PATH}/metadata_1000.csv"`. + +Optional variables: + +- `OUTPUT_CSV`: output CSV path, defaults to `${DATASET_BASE_PATH}/metadata_full.csv`. +- `SEGMENT_LENGTH`: frames per segment, default `81`. +- `CONTEXT_FRAMES`: context frames in metadata construction, default `5`. +- `METADATA_MAX_ROWS` / `DATASET_SIZE_ROWS`: optional row cap for custom-size metadata; `0` keeps the full CSV. + +## Latent precompute + +```bash +export WAN_BASE_MODEL=/path/to/Wan2.1-T2V-1.3B +export DATASET_BASE_PATH=data/Context-as-Memory-Dataset +NUM_PROCESSES=8 bash scripts/run_precompute_ctx_target_latents.sh +``` + +Optional variables: + +- `MODEL_PATHS`: JSON list of model weight paths. +- `CONTEXT_FRAMES`: number of context frames, default `20`. +- `NUM_PROCESSES`: accelerate processes, default `1`. + +Open-domain revisit first-frame assets are already included under `assets/opendomain_revisit`; no dataset construction step is needed for those probes. diff --git a/code/scripts/add_geometry_memory_column.py b/code/scripts/add_geometry_memory_column.py new file mode 100644 index 0000000000000000000000000000000000000000..e57d7416854b21bd0e6dd90b7c1261f57016827d --- /dev/null +++ b/code/scripts/add_geometry_memory_column.py @@ -0,0 +1,54 @@ +#!/usr/bin/env python3 +"""Attach pre-rendered static-geometry videos to an Echo-Memory metadata CSV.""" + +import argparse +from pathlib import Path + +import pandas as pd + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--metadata", required=True) + parser.add_argument("--geometry_root", required=True) + parser.add_argument("--output", required=True) + parser.add_argument("--column", default="geometry_memory") + parser.add_argument("--template", default="{video_name}/Vid_masktarget.mp4") + parser.add_argument("--allow_missing", action="store_true") + args = parser.parse_args() + + metadata_path = Path(args.metadata).resolve() + geometry_root = Path(args.geometry_root).resolve() + output_path = Path(args.output).resolve() + table = pd.read_csv(metadata_path) + if "video_name" not in table.columns: + raise ValueError("metadata must contain a video_name column") + + relative_paths = [] + missing = [] + for index, row in table.iterrows(): + values = {key: row[key] for key in table.columns} + relative = args.template.format(**values) + path = geometry_root / relative + if not path.is_file(): + missing.append((index, str(path))) + relative_paths.append(relative) + + if missing and not args.allow_missing: + examples = "\n".join(f" row {idx}: {path}" for idx, path in missing[:10]) + raise FileNotFoundError( + f"{len(missing)} geometry videos are missing under {geometry_root}. Examples:\n{examples}" + ) + + table[args.column] = relative_paths + output_path.parent.mkdir(parents=True, exist_ok=True) + table.to_csv(output_path, index=False) + print( + f"Wrote {len(table)} rows to {output_path}; " + f"geometry_column={args.column}, missing={len(missing)}" + ) + + +if __name__ == "__main__": + main() + diff --git a/code/scripts/precompute_ctx_target_latents.py b/code/scripts/precompute_ctx_target_latents.py new file mode 100644 index 0000000000000000000000000000000000000000..4521f6e2e8709cf8b7c1af09ba91faaf2b848e04 --- /dev/null +++ b/code/scripts/precompute_ctx_target_latents.py @@ -0,0 +1,546 @@ +#!/usr/bin/env python3 +""" +Precompute VAE latents for Context-as-Memory dataset with separate ctx and target storage: +- ctx: 1 latent per frame (for context/memory frames) +- target: 1 latent per 4 frames (time_division_factor=4) + +8-GPU distributed: each rank processes a subset of segments. + +Two modes: +1. With metadata: --metadata_path metadata_full.csv (uses VideoDataset) +2. No metadata: --no_metadata - auto-discovers segments from frames/ + captions.txt (or captions.jsonl) + +Usage (8 GPUs, no metadata): + accelerate launch --num_processes 8 scripts/precompute_ctx_target_latents.py \\ + --dataset_base_path /path/to/Context-as-Memory-Dataset \\ + --output_dir /path/to/latents \\ + --model_paths '["dit.safetensors","t5.pth","VAE.pth"]' \\ + --no_metadata +""" + +import argparse +import json +import os +import sys +import warnings +from datetime import datetime + +import torch +from PIL import Image +from tqdm import tqdm + +# Add project root for imports +SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +PROJECT_ROOT = os.path.abspath(os.path.join(SCRIPT_DIR, "..")) +if PROJECT_ROOT not in sys.path: + sys.path.insert(0, PROJECT_ROOT) + +from accelerate import Accelerator +from accelerate.utils import set_seed + +from diffsynth.pipelines.wan_video_new import WanVideoPipeline, ModelConfig +from diffsynth.trainers.utils import VideoDataset + + +def load_captions_txt(captions_path): + """Load captions.txt: video_name/start_end.mp4\\tcaption -> video_name -> caption.""" + captions = {} + if not os.path.isfile(captions_path): + return captions + with open(captions_path, "r", encoding="utf-8") as f: + for line in f: + parts = line.strip().split("\t", 1) + if len(parts) != 2: + continue + video_path, caption = parts[0], parts[1] + video_name = video_path.split("/")[0] + if video_name not in captions: + captions[video_name] = caption + return captions + + +def load_captions_jsonl(captions_path): + """Load captions.jsonl: each line {"video_name": "...", "prompt": "..."} or similar.""" + captions = {} + if not os.path.isfile(captions_path): + return captions + with open(captions_path, "r", encoding="utf-8") as f: + for line in f: + try: + obj = json.loads(line) + vn = obj.get("video_name") or obj.get("video") or obj.get("id", "") + prompt = obj.get("prompt") or obj.get("caption") or obj.get("text", "") + if vn: + captions[vn] = prompt + except json.JSONDecodeError: + continue + return captions + + +def build_segments_from_frames( + frames_dir, + captions_path=None, + captions_jsonl_path=None, + num_frames=81, + segment_stride=None, + overlap_labels_dir=None, + overlap_labels_dense=False, +): + """ + Build segment list without metadata CSV. + Modes: + 1. overlap_labels_dir + dense=False: 1 segment per overlap_labels JSON (~240k) + 2. overlap_labels_dir + dense=True: stride=1 within each video (~10x-40x more) + 3. else: stride-based from frames/ - stride 40 ~19k, stride 1 ~760k + + Returns: [(video_name, start_frame, end_frame, frame_paths, prompt), ...] + """ + captions = {} + if captions_path: + captions = load_captions_txt(captions_path) + if captions_jsonl_path: + captions.update(load_captions_jsonl(captions_jsonl_path)) + default_prompt = "A video scene." + + segments = [] + stride = segment_stride if segment_stride is not None else max(1, num_frames // 2) + + if overlap_labels_dir and os.path.isdir(overlap_labels_dir): + video_dirs = sorted( + [d for d in os.listdir(overlap_labels_dir) if os.path.isdir(os.path.join(overlap_labels_dir, d))] + ) + for video_name in video_dirs: + video_frames_dir = os.path.join(frames_dir, video_name) + if not os.path.isdir(video_frames_dir): + continue + frame_files = sorted([f for f in os.listdir(video_frames_dir) if f.endswith(".png")]) + if len(frame_files) < num_frames: + continue + prompt = captions.get(video_name, default_prompt) + + if overlap_labels_dense: + # Dense: stride=1 for max segments (~10x-40x more). Overridable via segment_stride. + seg_stride = segment_stride if segment_stride is not None else 1 + else: + # 1 segment per overlap_labels JSON + video_overlap_dir = os.path.join(overlap_labels_dir, video_name) + json_files = sorted([f for f in os.listdir(video_overlap_dir) if f.endswith(".json")]) + for jf in json_files: + try: + start_frame = int(jf.replace(".json", "")) + except ValueError: + continue + end_frame = start_frame + num_frames - 1 + frame_paths = [ + os.path.join(video_name, f"{start_frame + i:04d}.png") for i in range(num_frames) + ] + first_path = os.path.join(frames_dir, frame_paths[0]) + last_path = os.path.join(frames_dir, frame_paths[-1]) + if os.path.isfile(first_path) and os.path.isfile(last_path): + segments.append((video_name, start_frame, end_frame, frame_paths, prompt)) + continue + + for start in range(0, len(frame_files) - num_frames + 1, seg_stride): + end = start + num_frames - 1 + frame_paths = [os.path.join(video_name, frame_files[i]) for i in range(start, end + 1)] + segments.append((video_name, start, end, frame_paths, prompt)) + else: + # Stride-based from frames/ + video_dirs = sorted([d for d in os.listdir(frames_dir) if os.path.isdir(os.path.join(frames_dir, d))]) + for video_name in video_dirs: + video_dir = os.path.join(frames_dir, video_name) + frame_files = sorted([f for f in os.listdir(video_dir) if f.endswith(".png")]) + if len(frame_files) < num_frames: + continue + prompt = captions.get(video_name, default_prompt) + for start in range(0, len(frame_files) - num_frames + 1, stride): + end = start + num_frames - 1 + frame_paths = [os.path.join(video_name, frame_files[i]) for i in range(start, end + 1)] + segments.append((video_name, start, end, frame_paths, prompt)) + return segments + + +class FrameSegmentDataset(torch.utils.data.Dataset): + """Dataset that loads frames from segment list (no metadata CSV).""" + + def __init__(self, base_path, segments, height, width): + self.base_path = base_path + self.frames_dir = os.path.join(base_path, "frames") + self.segments = segments + self.height = height + self.width = width + + def __len__(self): + return len(self.segments) + + def _load_image(self, rel_path): + path = os.path.join(self.frames_dir, rel_path) + img = Image.open(path).convert("RGB") + import torchvision.transforms.functional as TF + w, h = img.size + scale = max(self.width / w, self.height / h) + img = TF.resize(img, (round(h * scale), round(w * scale)), interpolation=TF.InterpolationMode.BILINEAR) + img = TF.center_crop(img, (self.height, self.width)) + return img + + def __getitem__(self, idx): + video_name, start_frame, end_frame, frame_paths, prompt = self.segments[idx] + frames = [] + for fp in frame_paths: + try: + frames.append(self._load_image(fp)) + except Exception: + return None + if len(frames) != len(frame_paths): + return None + return { + "video": frames, + "prompt": prompt, + "video_name": video_name, + "start_frame": start_frame, + "end_frame": end_frame, + } + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Precompute ctx (1 latent/frame) and target (1 latent/4 frames) latents." + ) + parser.add_argument( + "--dataset_base_path", + type=str, + required=True, + help="Dataset root (contains frames/, metadata).", + ) + parser.add_argument( + "--metadata_path", + type=str, + default=None, + help="Metadata CSV path. Omit when using --no_metadata.", + ) + parser.add_argument( + "--no_metadata", + action="store_true", + help="Skip metadata CSV; auto-discover segments from frames/ + captions.", + ) + parser.add_argument( + "--captions_path", + type=str, + default=None, + help="captions.txt path (for --no_metadata). Default: {dataset_base_path}/captions.txt", + ) + parser.add_argument( + "--captions_jsonl_path", + type=str, + default=None, + help="Optional captions.jsonl path (for --no_metadata).", + ) + parser.add_argument( + "--segment_stride", + type=int, + default=None, + help="Stride between segments. Default: 40 (stride mode), 1 (--overlap_labels_dense). Use 1 for max.", + ) + parser.add_argument( + "--use_overlap_labels", + action="store_true", + help="When --no_metadata: use overlap_labels/ to discover segments (matches metadata_full ~240k).", + ) + parser.add_argument( + "--overlap_labels_dense", + action="store_true", + help="With --use_overlap_labels: stride=1 per video (~10x-40x more segments).", + ) + parser.add_argument( + "--output_dir", + type=str, + required=True, + help="Output directory. Will create ctx_latents/ and target_latents/ subdirs.", + ) + parser.add_argument( + "--model_paths", + type=str, + required=True, + help='JSON array of model paths, e.g. \'["dit.safetensors","t5.pth","Wan2.1_VAE.pth"]\'.', + ) + parser.add_argument("--tokenizer_path", type=str, default=None) + parser.add_argument("--height", type=int, default=352) + parser.add_argument("--width", type=int, default=640) + parser.add_argument("--num_frames", type=int, default=81) + parser.add_argument( + "--context_frames", + type=int, + default=5, + help="Number of context frames (each gets 1 latent).", + ) + parser.add_argument( + "--target_frames_per_latent", + type=int, + default=4, + help="Target: 1 latent per N frames (default 4).", + ) + parser.add_argument("--action_base_path", type=str, default=None) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--skip_existing", action="store_true") + parser.add_argument( + "--encode_batch_size", + type=int, + default=8, + help="Batch size for VAE encode (frames per call). Higher = better GPU util, more VRAM.", + ) + parser.add_argument( + "--segment_batch_size", + type=int, + default=8, + help="Process N segments per iteration (DataLoader batch). Speeds up I/O + encode.", + ) + return parser.parse_args() + + +def make_dataset_args(args): + return argparse.Namespace( + dataset_base_path=args.dataset_base_path, + dataset_metadata_path=args.metadata_path, + height=args.height, + width=args.width, + max_pixels=1920 * 1080, + num_frames=args.num_frames, + dataset_repeat=1, + data_file_keys="video,video_name,start_frame,end_frame", + action_base_path=args.action_base_path or args.dataset_base_path, + ) + + +def crop_and_resize(image, target_height, target_width): + import torchvision.transforms.functional as TF + width, height = image.size + scale = max(target_width / width, target_height / height) + image = TF.resize( + image, + (round(height * scale), round(width * scale)), + interpolation=TF.InterpolationMode.BILINEAR, + ) + image = TF.center_crop(image, (target_height, target_width)) + return image + + +def main(): + args = parse_args() + set_seed(args.seed) + + if not args.no_metadata and not args.metadata_path: + raise ValueError("Either --metadata_path or --no_metadata is required.") + + accelerator = Accelerator() + if accelerator.num_processes != 8 and accelerator.is_main_process: + print(f"Info: using {accelerator.num_processes} processes (expected 8).") + + ctx_dir = os.path.join(args.output_dir, "ctx_latents") + target_dir = os.path.join(args.output_dir, "target_latents") + os.makedirs(ctx_dir, exist_ok=True) + os.makedirs(target_dir, exist_ok=True) + + # Redirect rank 0 stdout to log file in output_dir + _rank0_log_file = None + if accelerator.is_main_process: + log_path = os.path.join(args.output_dir, "precompute_log.txt") + _rank0_log_file = open(log_path, "w", encoding="utf-8") + _rank0_log_file.write(f"[{datetime.now().isoformat()}] precompute_ctx_target_latents started\n") + _rank0_log_file.write(f"output_dir={args.output_dir}\n") + _rank0_log_file.flush() + + class _Tee: + def __init__(self, *files): + self.files = files + + def write(self, obj): + for f in self.files: + f.write(obj) + f.flush() + + def flush(self): + for f in self.files: + f.flush() + + sys.stdout = _Tee(sys.__stdout__, _rank0_log_file) + + if args.no_metadata: + frames_dir = os.path.join(args.dataset_base_path, "frames") + captions_path = args.captions_path or os.path.join(args.dataset_base_path, "captions.txt") + overlap_labels_dir = None + if args.use_overlap_labels: + overlap_labels_dir = os.path.join(args.dataset_base_path, "overlap_labels") + segments = build_segments_from_frames( + frames_dir, + captions_path=captions_path, + captions_jsonl_path=args.captions_jsonl_path, + num_frames=args.num_frames, + segment_stride=args.segment_stride, + overlap_labels_dir=overlap_labels_dir, + overlap_labels_dense=args.overlap_labels_dense, + ) + dataset = FrameSegmentDataset( + args.dataset_base_path, segments, args.height, args.width + ) + if accelerator.is_main_process: + if overlap_labels_dir: + src = "overlap_labels (dense)" if args.overlap_labels_dense else "overlap_labels" + else: + src = "frames (stride)" + print(f"No-metadata mode: discovered {len(segments)} segments from {src}") + else: + dataset_args = make_dataset_args(args) + dataset = VideoDataset(args=dataset_args) + + total = len(dataset) + if total == 0: + if accelerator.is_main_process: + print("Dataset is empty. Exit.") + return + + sampler = torch.utils.data.DistributedSampler( + dataset, + num_replicas=accelerator.num_processes, + rank=accelerator.process_index, + shuffle=False, + drop_last=False, + ) + indices = list(sampler) + n_local = len(indices) + + if accelerator.is_main_process: + print(f"Dataset size: {total}. Rank 0 processing {n_local} indices.") + meta = { + "dataset_base_path": args.dataset_base_path, + "metadata_path": args.metadata_path, + "no_metadata": args.no_metadata, + "total_samples": total, + "height": args.height, + "width": args.width, + "num_frames": args.num_frames, + "context_frames": args.context_frames, + "target_frames_per_latent": args.target_frames_per_latent, + } + with open(os.path.join(args.output_dir, "metadata_precompute.json"), "w") as f: + json.dump(meta, f, indent=2) + + # Load pipeline (VAE only) + model_paths = json.loads(args.model_paths) + model_configs = [ModelConfig(path=p) for p in model_paths] + from_pretrained_kw = { + "torch_dtype": torch.bfloat16, + "device": "cpu", + "model_configs": model_configs, + } + if args.tokenizer_path: + from_pretrained_kw["tokenizer_config"] = ModelConfig(path=args.tokenizer_path) + + if accelerator.is_main_process: + print("Loading pipeline (VAE)...") + pipe = WanVideoPipeline.from_pretrained(**from_pretrained_kw) + pipe.vae.to(accelerator.device) + pipe.vae.eval() + + K = args.context_frames + step = args.target_frames_per_latent + + def preprocess_frames(frames): + return pipe.preprocess_video(frames) + + @torch.no_grad() + def encode_frame(frame_pil): + """Encode single frame -> (C, 1, H//8, W//8)""" + vid = preprocess_frames([frame_pil]) + if vid.dim() == 5: + vid = vid.squeeze(0) + lat = pipe.vae.encode([vid], device=accelerator.device, tiled=False, tile_size=None, tile_stride=None) + return lat[0].cpu() + + failed = 0 + skipped = 0 + for idx in tqdm( + indices, + desc=f"Rank {accelerator.process_index}", + disable=not accelerator.is_local_main_process, + ): + sample = dataset[idx] + ctx_path = os.path.join(ctx_dir, f"{idx:08d}.pt") + target_path = os.path.join(target_dir, f"{idx:08d}.pt") + if args.skip_existing and os.path.isfile(ctx_path) and os.path.isfile(target_path): + skipped += 1 + continue + try: + if sample is None: + failed += 1 + continue + video_frames = sample.get("video") + if not video_frames or len(video_frames) == 0: + failed += 1 + continue + if len(video_frames) != args.num_frames: + if len(video_frames) > args.num_frames: + video_frames = video_frames[: args.num_frames] + else: + last = video_frames[-1] if video_frames else None + while len(video_frames) < args.num_frames and last is not None: + video_frames = video_frames + [last] + if len(video_frames) < args.num_frames: + failed += 1 + continue + + # Context: 1 latent per frame (frames 0..K-1) + ctx_latents_list = [] + for i in range(min(K, len(video_frames))): + lat = encode_frame(video_frames[i]) + if isinstance(lat, (list, tuple)): + lat = lat[0] + ctx_latents_list.append(lat) + ctx_latent = torch.cat(ctx_latents_list, dim=1) + if ctx_latent.dim() == 5: + ctx_latent = ctx_latent.squeeze(0) + + # Target: 1 latent per `step` frames (frames K, K+step, ...) + target_indices = list(range(K, len(video_frames), step)) + target_latents_list = [] + for i in target_indices: + lat = encode_frame(video_frames[i]) + if isinstance(lat, (list, tuple)): + lat = lat[0] + target_latents_list.append(lat) + if not target_latents_list: + failed += 1 + continue + target_latent = torch.cat(target_latents_list, dim=1) + if target_latent.dim() == 5: + target_latent = target_latent.squeeze(0) + + save_meta = { + "prompt": sample.get("prompt", ""), + "video_name": sample.get("video_name"), + "start_frame": sample.get("start_frame"), + "end_frame": sample.get("end_frame"), + } + if "actions" in sample and sample["actions"] is not None: + a = sample["actions"] + save_meta["actions"] = torch.tensor(a) if not isinstance(a, torch.Tensor) else a.cpu() + + torch.save({"latent": ctx_latent, **save_meta}, ctx_path) + torch.save({"latent": target_latent, **save_meta}, target_path) + except Exception as e: + if accelerator.is_local_main_process: + tqdm.write(f"Rank {accelerator.process_index} idx {idx}: {e}") + failed += 1 + + accelerator.wait_for_everyone() + if accelerator.is_main_process: + print( + f"Precompute done. ctx_latents/ and target_latents/ under {args.output_dir}. " + f"Failed: {failed}, Skipped: {skipped}." + ) + if _rank0_log_file is not None: + sys.stdout = sys.__stdout__ + _rank0_log_file.close() + + +if __name__ == "__main__": + main() diff --git a/code/scripts/publish_gh_pages.sh b/code/scripts/publish_gh_pages.sh new file mode 100644 index 0000000000000000000000000000000000000000..9ad60eb5a796b88fb38ea7efd8b15c046cfcd44e --- /dev/null +++ b/code/scripts/publish_gh_pages.sh @@ -0,0 +1,49 @@ +#!/usr/bin/env bash +# Publish docs/ to the gh-pages branch (GitHub Pages source). +set -euo pipefail + +ROOT="$(cd "$(dirname "$0")/.." && pwd)" +cd "$ROOT" + +if ! git rev-parse --is-inside-work-tree >/dev/null 2>&1; then + echo "error: not inside a git repository" >&2 + exit 1 +fi + +if [[ ! -f docs/index.html || ! -f docs/style.css || ! -f docs/site.js || ! -f docs/i18n.js ]]; then + echo "error: docs/index.html, docs/style.css, docs/site.js, and docs/i18n.js are required" >&2 + exit 1 +fi + +MAIN_SHA="$(git rev-parse --short HEAD)" +WORKTREE="$(mktemp -d)" +trap 'git worktree remove -f "$WORKTREE" 2>/dev/null || true; rm -rf "$WORKTREE"' EXIT + +git fetch origin gh-pages 2>/dev/null || true +if git show-ref --verify --quiet refs/remotes/origin/gh-pages; then + git worktree add -B gh-pages-publish "$WORKTREE" origin/gh-pages +else + git worktree add -B gh-pages-publish "$WORKTREE" --orphan gh-pages-publish +fi + +find "$WORKTREE" -mindepth 1 -maxdepth 1 ! -name '.git' -exec rm -rf {} + + +cp docs/index.html docs/style.css docs/site.js docs/i18n.js docs/i18n-runtime.js docs/developer.html "$WORKTREE/" +cp docs/.nojekyll "$WORKTREE/" 2>/dev/null || : >"$WORKTREE/.nojekyll" +rm -rf "$WORKTREE/assets" +mkdir -p "$WORKTREE/assets" +cp -a docs/assets/. "$WORKTREE/assets/" + +# Cache-bust marker for verifying deploys in page source. +sed -i "s/site-build-main/site-build-${MAIN_SHA}/" "$WORKTREE/index.html" + +cd "$WORKTREE" +git add -A +if git diff --cached --quiet; then + echo "gh-pages: no changes to publish" + exit 0 +fi + +git commit -m "Publish project page from main (${MAIN_SHA})" +git push origin HEAD:gh-pages +echo "Published docs/ to origin/gh-pages (${MAIN_SHA})" diff --git a/code/scripts/run_generate_metadata.sh b/code/scripts/run_generate_metadata.sh new file mode 100644 index 0000000000000000000000000000000000000000..bdbd4b01b834b6a6ab7da487ccac585a69d193d7 --- /dev/null +++ b/code/scripts/run_generate_metadata.sh @@ -0,0 +1,70 @@ +#!/bin/bash +# Generate metadata CSV for a context-based memory dataset. +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="${REPO_ROOT:-$(cd "${SCRIPT_DIR}/.." && pwd)}" +cd "${REPO_ROOT}" + +DETECTED_CPUS=$(python3 -c "import os; print(os.cpu_count())") +OPTIMAL_WORKERS=$((DETECTED_CPUS > 4 ? DETECTED_CPUS - 2 : DETECTED_CPUS)) + +NUM_WORKERS=${NUM_WORKERS:-$OPTIMAL_WORKERS} +DATASET_BASE_PATH="${DATASET_BASE_PATH:-${REPO_ROOT}/data/Context-as-Memory-Dataset}" +OUTPUT_CSV="${OUTPUT_CSV:-${DATASET_BASE_PATH}/metadata_full.csv}" +SEGMENT_LENGTH="${SEGMENT_LENGTH:-81}" +CONTEXT_FRAMES="${CONTEXT_FRAMES:-5}" +METADATA_MAX_ROWS="${METADATA_MAX_ROWS:-${DATASET_SIZE_ROWS:-0}}" + +echo "==========================================" +echo "Context-as-Memory Metadata Generation" +echo "==========================================" +echo "Detected CPUs: $DETECTED_CPUS" +echo "Using workers: $NUM_WORKERS" +echo "Dataset: ${DATASET_BASE_PATH}" +echo "Output: ${OUTPUT_CSV}" +echo "Max rows: ${METADATA_MAX_ROWS} (0 = full metadata)" +echo "==========================================" +echo "" + +python3 src/data/preprocess_cam_dataset.py \ + --dataset_base_path "${DATASET_BASE_PATH}" \ + --output_csv "${OUTPUT_CSV}" \ + --segment_length "${SEGMENT_LENGTH}" \ + --context_frames "${CONTEXT_FRAMES}" + +if [ "${METADATA_MAX_ROWS}" != "0" ]; then + python3 - "${OUTPUT_CSV}" "${METADATA_MAX_ROWS}" <<'PY' +import csv +import os +import sys + +path = sys.argv[1] +max_rows = int(sys.argv[2]) +if max_rows < 0: + raise SystemExit("METADATA_MAX_ROWS must be >= 0") +if max_rows > 0: + tmp_path = f"{path}.tmp" + with open(path, newline="", encoding="utf-8") as src, open(tmp_path, "w", newline="", encoding="utf-8") as dst: + reader = csv.reader(src) + writer = csv.writer(dst) + header = next(reader, None) + if header is not None: + writer.writerow(header) + for idx, row in enumerate(reader): + if idx >= max_rows: + break + writer.writerow(row) + os.replace(tmp_path, path) + print(f"Truncated metadata to {max_rows} rows: {path}") +PY +fi + +echo "" +echo "Generation complete! Checking results..." +if [ -f "${OUTPUT_CSV}" ]; then + echo "CSV file generated successfully" + wc -l "${OUTPUT_CSV}" +else + echo "CSV file generation failed" +fi diff --git a/code/scripts/run_precompute_ctx_target_latents.sh b/code/scripts/run_precompute_ctx_target_latents.sh new file mode 100644 index 0000000000000000000000000000000000000000..c70cb4d0a0507a6c2e3a8186568b553afba204a2 --- /dev/null +++ b/code/scripts/run_precompute_ctx_target_latents.sh @@ -0,0 +1,32 @@ +#!/bin/bash +# Precompute ctx (1 latent/frame) and target (1 latent/4 frames) latents for Context-as-Memory dataset. +# +# No metadata needed: auto-discovers segments from frames/ + captions.txt + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="${REPO_ROOT:-$(cd "${SCRIPT_DIR}/.." && pwd)}" + +DATASET_BASE="${DATASET_BASE_PATH:-${REPO_ROOT}/data/Context-as-Memory-Dataset}" +OUTPUT_DIR="${DATASET_BASE}/latents" +WAN_BASE_MODEL="${WAN_BASE_MODEL:-${REPO_ROOT}/checkpoints/Wan2.1-T2V-1.3B}" +MODEL_PATHS="${MODEL_PATHS:-[\"${WAN_BASE_MODEL}/diffusion_pytorch_model.safetensors\",\"${WAN_BASE_MODEL}/models_t5_umt5-xxl-enc-bf16.pth\",\"${WAN_BASE_MODEL}/Wan2.1_VAE.pth\"]}" +NUM_PROCESSES="${NUM_PROCESSES:-1}" +CONTEXT_FRAMES="${CONTEXT_FRAMES:-20}" + +cd "${REPO_ROOT}" + +accelerate launch --num_processes "${NUM_PROCESSES}" scripts/precompute_ctx_target_latents.py \ + --dataset_base_path "${DATASET_BASE}" \ + --output_dir "${OUTPUT_DIR}" \ + --model_paths "${MODEL_PATHS}" \ + --height 352 --width 640 \ + --context_frames "${CONTEXT_FRAMES}" \ + --target_frames_per_latent 4 \ + --no_metadata \ + --use_overlap_labels \ + --overlap_labels_dense \ + --skip_existing + +echo "Done. Latents saved to ${OUTPUT_DIR}/ctx_latents/ and ${OUTPUT_DIR}/target_latents/" diff --git a/code/scripts/run_spmem_tsdf_preprocess.sh b/code/scripts/run_spmem_tsdf_preprocess.sh new file mode 100644 index 0000000000000000000000000000000000000000..ca24e52f6c20970406f9e91e01ea10c439e015d4 --- /dev/null +++ b/code/scripts/run_spmem_tsdf_preprocess.sh @@ -0,0 +1,33 @@ +#!/bin/bash +# Run the official arXiv:2506.05284 TSDF renderer on one reconstructed clip. +# INPUT_NPZ must contain: images, depths, intrinsic, cam_c2w. +set -euo pipefail + +SPMEM_ROOT="${SPMEM_ROOT:?Set SPMEM_ROOT to a checkout of https://github.com/spmem/spmem}" +INPUT_NPZ="${INPUT_NPZ:?Set INPUT_NPZ to the reconstructed clip .npz}" +CLIP_NAME="${CLIP_NAME:?Set CLIP_NAME to the output sample name}" +GEOMETRY_OUTPUT_ROOT="${GEOMETRY_OUTPUT_ROOT:?Set GEOMETRY_OUTPUT_ROOT}" + +RUN_DATA="${SPMEM_ROOT}/tsdf/run_data.py" +if [ ! -f "${RUN_DATA}" ]; then + echo "Missing official TSDF entrypoint: ${RUN_DATA}" >&2 + exit 2 +fi +if [ ! -f "${INPUT_NPZ}" ]; then + echo "Missing reconstructed clip: ${INPUT_NPZ}" >&2 + exit 2 +fi + +python "${RUN_DATA}" \ + --input_dir "${INPUT_NPZ}" \ + --name "${CLIP_NAME}" \ + --save_dir "${GEOMETRY_OUTPUT_ROOT}" \ + --fast_mode + +OUTPUT_VIDEO="${GEOMETRY_OUTPUT_ROOT}/${CLIP_NAME}/Vid_masktarget.mp4" +if [ ! -f "${OUTPUT_VIDEO}" ]; then + echo "TSDF preprocessing completed without expected output: ${OUTPUT_VIDEO}" >&2 + exit 3 +fi +echo "${OUTPUT_VIDEO}" + diff --git a/code/src/data/__init__.py b/code/src/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..01c4482392f81ea64352569834bf3914a523d21e --- /dev/null +++ b/code/src/data/__init__.py @@ -0,0 +1 @@ +"""Data preprocessing utilities for Echo-Memory.""" diff --git a/code/src/data/preprocess_cam_dataset.py b/code/src/data/preprocess_cam_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..3c2232b95567ec825e77298ac9eacb6e80a046cc --- /dev/null +++ b/code/src/data/preprocess_cam_dataset.py @@ -0,0 +1,232 @@ +#!/usr/bin/env python3 +""" +Preprocess Context-as-Memory dataset folders into Echo-Memory metadata CSV. + +Expected dataset layout: +- frames/: frame images organized by video +- jsons/: camera pose information for each video +- overlap_labels/: FOV overlap information for memory retrieval +- captions.txt: video segment captions +""" + +import argparse +import csv +import json +import os +from typing import Dict, List, Tuple + + +def parse_caption_line(line: str) -> Tuple[str, str]: + """ + Parse a line from captions.txt. + + Format: "video_name/start_end.mp4\tcaption text..." + Returns: (video_path, caption) + """ + parts = line.strip().split("\t", 1) + if len(parts) != 2: + return None, None + video_path = parts[0] + caption = parts[1] + return video_path, caption + + +def load_captions(captions_file: str) -> Dict[str, str]: + """Load captions.txt as video_name -> caption.""" + captions = {} + if not os.path.exists(captions_file): + print(f"Warning: Captions file not found: {captions_file}") + return captions + + with open(captions_file, "r", encoding="utf-8") as f: + for line in f: + video_path, caption = parse_caption_line(line) + if video_path and caption: + video_name = video_path.split("/")[0] + if video_name not in captions: + captions[video_name] = [] + captions[video_name].append(caption) + + for video_name in captions: + captions[video_name] = captions[video_name][0] if captions[video_name] else "" + + return captions + + +def get_frame_files(frames_dir: str, video_name: str) -> List[str]: + """Get sorted frame paths for one video, relative to frames_dir.""" + video_frames_dir = os.path.join(frames_dir, video_name) + if not os.path.exists(video_frames_dir): + return [] + + frame_files = [] + for frame_file in sorted(os.listdir(video_frames_dir)): + if frame_file.endswith(".png"): + frame_files.append(os.path.join(video_name, frame_file)) + + return frame_files + + +def load_camera_poses(json_file: str) -> Dict: + """Load camera poses from a JSON file.""" + if not os.path.exists(json_file): + return {} + + with open(json_file, "r", encoding="utf-8") as f: + data = json.load(f) + + if "CineCameraActor" in data: + return data["CineCameraActor"] + if isinstance(data, dict): + return data + return {} + + +def load_overlap_labels(overlap_dir: str, video_name: str, frame_idx: int) -> List[int]: + """Load overlapping frame indices for a given frame.""" + overlap_file = os.path.join(overlap_dir, video_name, f"{frame_idx}.json") + if not os.path.exists(overlap_file): + return [] + + try: + with open(overlap_file, "r", encoding="utf-8") as f: + data = json.load(f) + overlapping_frames = data.get("overlapping_frames", []) + return [int(frame) for frame in overlapping_frames if str(frame).isdigit()] + except Exception: + return [] + + +def create_metadata_csv( + dataset_base_path: str, + output_csv: str, + segment_length: int = 81, + context_frames: int = 5, +): + """ + Create metadata CSV for the Context-as-Memory dataset. + + Args: + dataset_base_path: root of the dataset. + output_csv: output CSV path. + segment_length: frames per training segment. + context_frames: context frames reserved by downstream workflows. + """ + frames_dir = os.path.join(dataset_base_path, "frames") + captions_file = os.path.join(dataset_base_path, "captions.txt") + + captions = load_captions(captions_file) + + if not os.path.exists(frames_dir): + print(f"Error: Frames directory not found: {frames_dir}") + return + + video_names = [ + d for d in os.listdir(frames_dir) + if os.path.isdir(os.path.join(frames_dir, d)) + ] + + print(f"Found {len(video_names)} videos") + print(f"Context frames: {context_frames}") + + output_dir = os.path.dirname(output_csv) + if output_dir: + os.makedirs(output_dir, exist_ok=True) + + with open(output_csv, "w", newline="", encoding="utf-8") as csvfile: + fieldnames = [ + "video", + "prompt", + "video_name", + "start_frame", + "end_frame", + ] + writer = csv.DictWriter(csvfile, fieldnames=fieldnames) + writer.writeheader() + + total_segments = 0 + + for video_name in sorted(video_names): + print(f"Processing video: {video_name}") + + frame_files = get_frame_files(frames_dir, video_name) + if len(frame_files) < segment_length: + print( + f" Skipping {video_name}: only {len(frame_files)} frames " + f"(need at least {segment_length})" + ) + continue + + prompt = captions.get(video_name, f"A scene from {video_name}") + step = max(1, segment_length // 2) + video_segments = 0 + + for start_idx in range(0, len(frame_files) - segment_length + 1, step): + end_idx = start_idx + segment_length - 1 + segment_frames = frame_files[start_idx:end_idx + 1] + + if len(segment_frames) < segment_length: + continue + + frame_paths = "|".join(segment_frames) + video_path = os.path.join("frames", frame_paths) + + writer.writerow({ + "video": video_path, + "prompt": prompt, + "video_name": video_name, + "start_frame": start_idx, + "end_frame": end_idx, + }) + + total_segments += 1 + video_segments += 1 + + print(f" Created {video_segments} segments for {video_name}") + + print(f"\nTotal segments created: {total_segments}") + print(f"Metadata CSV saved to: {output_csv}") + + +def main(): + parser = argparse.ArgumentParser(description="Preprocess Context-as-Memory Dataset") + parser.add_argument( + "--dataset_base_path", + type=str, + required=True, + help="Base path to Context-as-Memory dataset", + ) + parser.add_argument( + "--output_csv", + type=str, + default="metadata.csv", + help="Output CSV file path (default: metadata.csv)", + ) + parser.add_argument( + "--segment_length", + type=int, + default=81, + help="Length of video segments (default: 81 frames)", + ) + parser.add_argument( + "--context_frames", + type=int, + default=5, + help="Number of context frames (default: 5)", + ) + + args = parser.parse_args() + + if not os.path.isabs(args.output_csv): + args.output_csv = os.path.join(args.dataset_base_path, args.output_csv) + + create_metadata_csv( + dataset_base_path=args.dataset_base_path, + output_csv=args.output_csv, + segment_length=args.segment_length, + context_frames=args.context_frames, + ) + + +if __name__ == "__main__": + main() diff --git a/code/src/model_training/fov_retrieval.py b/code/src/model_training/fov_retrieval.py new file mode 100644 index 0000000000000000000000000000000000000000..d01366081fa0f2c3887a2df6f446d80eb6e7ed38 --- /dev/null +++ b/code/src/model_training/fov_retrieval.py @@ -0,0 +1,1509 @@ +""" +FOV (Field of View) Overlap-based Memory Retrieval Module +Implements geometric retrieval based on camera pose overlap for Context-as-Memory + +Aligned with the Context-as-Memory paper [2506.03141]. +""" + +import os +import json +import random +import sys +from concurrent.futures import ThreadPoolExecutor, as_completed +import numpy as np +from typing import List, Dict, Tuple, Optional +from pathlib import Path + +import torch +from PIL import Image + +def degrees_to_radians(degrees: float) -> float: + """Convert degrees to radians.""" + return degrees * np.pi / 180 + + +def compute_rotation_list_z_only(x: float, y: float, z: float, yaw_degrees: float) -> List[float]: + """ + Paper-aligned: rotation only around Z-axis (yaw). 4 parameters. + R = R_z(yaw). No roll/pitch. + + Returns: + [t_x, t_y, t_z, R_11, R_12, ..., R_33] (12 elements) + """ + yaw_rad = degrees_to_radians(yaw_degrees) + c, s = np.cos(yaw_rad), np.sin(yaw_rad) + R_z = np.array([ + [c, -s, 0], + [s, c, 0], + [0, 0, 1], + ], dtype=np.float64) + return [float(x), float(y), float(z)] + R_z.flatten().tolist() + + +def compute_rotation_list_yaw_pitch( + x: float, y: float, z: float, yaw_degrees: float, pitch_degrees: float +) -> List[float]: + """ + Five-parameter rotation: R = R_z(pitch) @ R_y(yaw). Same convention as external pipelines. + - R_y: rotation around Y (yaw in degrees). + - R_z: rotation around Z (pitch in degrees). + Returns [t_x, t_y, t_z, R_11, ..., R_33] (12 elements), encoder-compatible. + """ + yaw_rad = degrees_to_radians(yaw_degrees) + pitch_rad = degrees_to_radians(pitch_degrees) + cy, sy = np.cos(yaw_rad), np.sin(yaw_rad) + cp, sp = np.cos(pitch_rad), np.sin(pitch_rad) + R_y = np.array([ + [cy, 0, sy], + [0, 1, 0], + [-sy, 0, cy], + ], dtype=np.float64) + R_z = np.array([ + [cp, -sp, 0], + [sp, cp, 0], + [0, 0, 1], + ], dtype=np.float64) + R = R_z @ R_y + return [float(x), float(y), float(z)] + R.flatten().tolist() + + +def compute_rotation_list(params: List[float]) -> List[float]: + """ + Compute rotation matrix from camera parameters and flatten to list. + - 4 params [x, y, z, yaw]: Z-only, R = R_z(yaw) (paper default). + - 5 params [x, y, z, yaw, pitch]: R = R_z(pitch) @ R_y(yaw). + Returns [t_x, t_y, t_z, R_11, ..., R_33] (12 elements). + """ + if len(params) >= 5: + x, y, z = float(params[0]), float(params[1]), float(params[2]) + yaw = float(params[3]) + pitch = float(params[4]) + return compute_rotation_list_yaw_pitch(x, y, z, yaw, pitch) + if len(params) >= 4: + x, y, z = float(params[0]), float(params[1]), float(params[2]) + yaw = float(params[3]) + return compute_rotation_list_z_only(x, y, z, yaw) + x = y = z = yaw = 0.0 + return compute_rotation_list_z_only(x, y, z, yaw) + + +def yaw_deg_from_rt(rt: List[float]) -> float: + """Extract yaw (degrees) from 12-dim RT. Z-only: yaw = atan2(R_21, R_11).""" + if rt is None or len(rt) < 12: + return 0.0 + R = np.array(rt[3:12]).reshape(3, 3) + return float(np.degrees(np.arctan2(R[1, 0], R[0, 0]))) + + +def flip_yaw_rt(rt: List[float]) -> List[float]: + """Return RT with yaw negated (CW<->CCW). rt is 12 elements [t_x,t_y,t_z, R_11..R_33].""" + if rt is None or len(rt) < 12: + return list(rt) if rt else [] + tx, ty, tz = rt[0], rt[1], rt[2] + yaw_deg = yaw_deg_from_rt(rt) + return compute_rotation_list_z_only(tx, ty, tz, -yaw_deg) + + +def flip_yaw_rt_list(rt_list: List[List[float]]) -> List[List[float]]: + """Flip yaw for each RT in list (data aug for direction sensitivity).""" + return [flip_yaw_rt(rt) for rt in rt_list] + + +def convert_rt_to_relative(rt_list_all: List[List[float]], ref_rt: List[float]) -> List[List[float]]: + """ + Convert RT (rotation-translation) poses to relative coordinates. + + This aligns with the Context-as-Memory paper's use of relative camera poses + for better geometric consistency in context frame retrieval. + + Args: + rt_list_all: List of RT poses, each is [t_x, t_y, t_z, R_11, R_12, ..., R_33] (12 elements) + ref_rt: Reference RT pose [t_x, t_y, t_z, R_11, R_12, ..., R_33] (12 elements) + + Returns: + new_rt_list: List of relative RT poses in the same format + """ + def parse_rt(rt: List[float]) -> tuple: + """Parse RT list into rotation matrix R and translation vector t.""" + t = np.array(rt[:3]).reshape((3, 1)) + R = np.array(rt[3:]).reshape((3, 3)) + return R, t + + R_ref, T_ref = parse_rt(ref_rt) + R_ref_inv = R_ref.T + T_ref_inv = -R_ref_inv @ T_ref + + new_rt_list = [] + + for rt in rt_list_all: + R_i, T_i = parse_rt(rt) + + # Convert to relative coordinates + R_new = R_ref_inv @ R_i + T_new = R_ref_inv @ T_i + T_ref_inv + + # Flatten back to list format: [t_x, t_y, t_z, R_11, R_12, ..., R_33] + rt_new = T_new.flatten().tolist() + R_new.flatten().tolist() + new_rt_list.append(rt_new) + + return new_rt_list + + +def pose_to_rt(pose: Dict, constrain_to_xy: bool = True) -> Optional[List[float]]: + """ + Convert camera pose dict to RT format. Paper-aligned by default. + + - 2D plane (constrain_to_xy=True): use position x, y only; z=0. + - Rotation: Z-axis only; use rotation[2] as yaw (degrees). Roll/pitch ignored. + + Args: + pose: Dict with 'position' [x, y, z] and 'rotation' [roll, pitch, yaw] in degrees + constrain_to_xy: If True (default), set z=0 for strict XY-plane displacement (paper). + + Returns: + RT list: [t_x, t_y, t_z, R_11, R_12, ..., R_33] (12 elements) or None if invalid + """ + if pose is None: + return None + + pos = pose.get('position', [0, 0, 0]) + rot = pose.get('rotation', [0, 0, 0]) + + if len(pos) < 2: + return None + + x = float(pos[0]) + y = float(pos[1]) + z = 0.0 if constrain_to_xy else (float(pos[2]) if len(pos) >= 3 else 0.0) + # Paper: rotation only around Z-axis -> use yaw (index 2) only + yaw = float(rot[2]) if len(rot) > 2 else 0.0 + + return compute_rotation_list([x, y, z, yaw]) + + +def rt_to_pose(rt: List[float]) -> Optional[Dict]: + """ + Convert RT format back to pose dict. + + Args: + rt: RT list [t_x, t_y, t_z, R_11, R_12, ..., R_33] (12 elements) + + Returns: + pose dict with 'position' and 'rotation' or None if invalid + """ + if rt is None or len(rt) < 12: + return None + + t = np.array(rt[:3]) + R = np.array(rt[3:]).reshape((3, 3)) + + # Extract Euler angles from rotation matrix + # Using ZYX convention (yaw-pitch-roll) + sy = np.sqrt(R[0, 0] * R[0, 0] + R[1, 0] * R[1, 0]) + singular = sy < 1e-6 + + if not singular: + roll = np.arctan2(R[2, 1], R[2, 2]) + pitch = np.arctan2(-R[2, 0], sy) + yaw = np.arctan2(R[1, 0], R[0, 0]) + else: + roll = np.arctan2(-R[1, 2], R[1, 1]) + pitch = np.arctan2(-R[2, 0], sy) + yaw = 0 + + return { + 'position': t.tolist(), + 'rotation': [np.degrees(roll), np.degrees(pitch), np.degrees(yaw)] + } + + + +def _parse_poses_dict(data: dict) -> dict: + """Extract poses dict from JSON data (CineCameraActor or flat dict).""" + if 'CineCameraActor' in data: + return data['CineCameraActor'] + return data if isinstance(data, dict) else {} + + +def load_poses_dict(json_file: str) -> dict: + """ + Load full camera poses dict from JSON file (one read for all frames). + + Args: + json_file: Path to camera pose JSON file + + Returns: + Dict mapping frame_idx (str) -> pose dict, or {} if failed + """ + if not os.path.exists(json_file): + return {} + try: + with open(json_file, 'r') as f: + data = json.load(f) + return _parse_poses_dict(data) + except Exception as e: + print(f"Error loading poses from {json_file}: {e}") + return {} + + +def load_camera_pose(json_file: str, frame_idx: int) -> Optional[Dict]: + """ + Load camera pose for a specific frame from JSON file. + + Args: + json_file: Path to camera pose JSON file + frame_idx: Frame index + + Returns: + Dict with 'position' and 'rotation' keys, or None if not found + """ + poses = load_poses_dict(json_file) + frame_key = str(frame_idx) + return poses.get(frame_key) + + +def load_camera_poses_batch(json_file: str, frame_indices: List[int]) -> List[Optional[Dict]]: + """ + Load camera poses for multiple frames in one JSON read. + + Args: + json_file: Path to camera pose JSON file + frame_indices: List of frame indices + + Returns: + List of pose dicts (or None) in same order as frame_indices + """ + poses = load_poses_dict(json_file) + return [poses.get(str(fi)) for fi in frame_indices] + + +def load_overlap_frames(overlap_labels_dir: str, video_name: str, frame_idx: int) -> List[int]: + """ + Load overlapping frame indices for a given frame from overlap_labels. + + Args: + overlap_labels_dir: Base directory for overlap labels + video_name: Name of the video + frame_idx: Current frame index + + Returns: + List of overlapping frame indices + """ + overlap_file = os.path.join(overlap_labels_dir, video_name, f"{frame_idx}.json") + if not os.path.exists(overlap_file): + return [] + + try: + # Add distributed training safety - timeout protection + import signal + import torch.distributed as dist + + def timeout_handler(signum, frame): + raise TimeoutError(f"Timeout loading overlap file: {overlap_file}") + + # Set 10-second timeout for file operations in distributed training + if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1: + signal.signal(signal.SIGALRM, timeout_handler) + signal.alarm(10) + + try: + with open(overlap_file, 'r') as f: + data = json.load(f) + overlapping_frames = data.get('overlapping_frames', []) + # Convert string indices to integers + result = [int(f) for f in overlapping_frames if f.isdigit() or isinstance(f, int)] + if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1: + signal.alarm(0) # Cancel alarm + return result + finally: + if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1: + signal.alarm(0) # Always cancel alarm + + except TimeoutError as e: + print(f"Warning: Timeout loading overlap file {overlap_file} (distributed training): {e}") + return [] + except Exception as e: + print(f"Error loading overlap labels from {overlap_file}: {e}") + return [] + + +def compute_fov_overlap_3d( + pose1: Dict, + pose2: Dict, + fov_degrees: float = 52.67, + max_distance: float = 500.0 # Increased default from 50.0 to 500.0 for large scenes +) -> float: + """ + Compute FOV overlap score between two camera poses using 3D geometry. + + This implementation uses full 6-DoF camera poses (position + rotation) + to compute more accurate FOV overlap, as described in Context-as-Memory. + + Args: + pose1: First camera pose with 'position' [x, y, z] and 'rotation' [roll, pitch, yaw] in degrees + pose2: Second camera pose with 'position' [x, y, z] and 'rotation' [roll, pitch, yaw] in degrees + fov_degrees: Field of view in degrees (default: 52.67 from paper) + max_distance: Maximum distance to consider (meters) + + Returns: + Overlap score between 0 and 1 + """ + if pose1 is None or pose2 is None: + return 0.0 + + pos1 = np.array(pose1.get('position', [0, 0, 0]), dtype=np.float32) + pos2 = np.array(pose2.get('position', [0, 0, 0]), dtype=np.float32) + + # Distance between cameras + distance = np.linalg.norm(pos2 - pos1) + # Instead of returning 0.0 for distances > max_distance, we use a soft threshold + # that still gives some score based on direction similarity even for far cameras + distance_exceeds_max = distance > max_distance + + if distance < 1e-6: + # Same position - high overlap + return 1.0 + + # Extract rotation angles (assuming [roll, pitch, yaw] or [x, y, z] rotation in degrees) + rot1 = pose1.get('rotation', [0, 0, 0]) + rot2 = pose2.get('rotation', [0, 0, 0]) + + # Convert to numpy array + rot1 = np.array(rot1, dtype=np.float32) + rot2 = np.array(rot2, dtype=np.float32) + + # Compute rotation matrices from Euler angles + # Note: The rotation order may vary by dataset. Common conventions: + # - ZYX (yaw-pitch-roll): R = R_z(yaw) * R_y(pitch) * R_x(roll) + # - XYZ (roll-pitch-yaw): R = R_x(roll) * R_y(pitch) * R_z(yaw) + # Based on Context-as-Memory dataset, rotation[2] is yaw (rotation around Z-axis) + # We'll use ZYX convention: yaw (Z), pitch (Y), roll (X) + + def euler_to_rotation_matrix(euler_angles): + """Convert Euler angles [roll, pitch, yaw] in degrees to rotation matrix (ZYX order)""" + roll, pitch, yaw = np.radians(euler_angles) + + # Rotation around X-axis (roll) + Rx = np.array([ + [1, 0, 0], + [0, np.cos(roll), -np.sin(roll)], + [0, np.sin(roll), np.cos(roll)] + ]) + + # Rotation around Y-axis (pitch) + Ry = np.array([ + [np.cos(pitch), 0, np.sin(pitch)], + [0, 1, 0], + [-np.sin(pitch), 0, np.cos(pitch)] + ]) + + # Rotation around Z-axis (yaw) + Rz = np.array([ + [np.cos(yaw), -np.sin(yaw), 0], + [np.sin(yaw), np.cos(yaw), 0], + [0, 0, 1] + ]) + + # ZYX order: R = Rz * Ry * Rx + R = Rz @ Ry @ Rx + return R + + # Handle different rotation formats + if len(rot1) >= 3: + # Euler angles [roll, pitch, yaw] or [x, y, z] + R1 = euler_to_rotation_matrix([rot1[0], rot1[1], rot1[2]]) + elif len(rot1) == 9: + # Rotation matrix flattened (3x3 = 9 elements) + R1 = rot1.reshape(3, 3) + else: + # Fallback: only yaw + R1 = euler_to_rotation_matrix([0, 0, rot1[2] if len(rot1) > 2 else 0]) + + if len(rot2) >= 3: + R2 = euler_to_rotation_matrix([rot2[0], rot2[1], rot2[2]]) + elif len(rot2) == 9: + R2 = rot2.reshape(3, 3) + else: + R2 = euler_to_rotation_matrix([0, 0, rot2[2] if len(rot2) > 2 else 0]) + + # Camera forward vector (typically Z-axis in camera coordinate system) + # In OpenCV/OpenGL convention, forward is usually -Z or +Z + # Based on Context-as-Memory dataset, we assume forward is +Z (third column) + forward1 = R1[:, 2] # Third column of rotation matrix + forward2 = R2[:, 2] + + # Vector from camera1 to camera2 + vec_1_to_2 = pos2 - pos1 + vec_1_to_2_norm = np.linalg.norm(vec_1_to_2) + vec_1_to_2_unit = vec_1_to_2 / (vec_1_to_2_norm + 1e-6) + + # FOV half-angle threshold (cosine of half FOV) + fov_rad = np.radians(fov_degrees) + fov_half_cos = np.cos(fov_rad / 2) + + # Check if camera1 can see camera2's position (within FOV) + # cos(angle) = dot(forward, vec_to_target) + # angle < fov/2 => cos(angle) > cos(fov/2) + dot1 = np.dot(forward1, vec_1_to_2_unit) + can_1_see_2 = dot1 > fov_half_cos + + # Check if camera2 can see camera1's position (within FOV) + dot2 = np.dot(forward2, -vec_1_to_2_unit) # Negative because looking back + can_2_see_1 = dot2 > fov_half_cos + + # Compute overlap score based on mutual visibility and distance + # Normalize distance for scoring (use max_distance as reference, but don't hard-cut) + normalized_distance = min(1.0, distance / max_distance) if max_distance > 0 else 1.0 + + if can_1_see_2 and can_2_see_1: + # Both cameras can see each other - high overlap + # Score decreases with distance, but never goes to 0 + distance_factor = 1.0 - normalized_distance * 0.5 + overlap = 0.8 + 0.2 * distance_factor + elif can_1_see_2 or can_2_see_1: + # One camera can see the other - medium overlap + distance_factor = 1.0 - normalized_distance * 0.6 + overlap = 0.4 + 0.3 * distance_factor + else: + # Check if cameras are looking in similar directions (even if not directly at each other) + # This handles the case where both cameras see the same scene from different angles + forward_similarity = np.dot(forward1, forward2) + if forward_similarity > 0.7: # Cameras looking in similar directions + # Even for far cameras, if they're looking in similar directions, there's some overlap + distance_factor = 1.0 - normalized_distance * 0.7 + overlap = 0.2 + 0.3 * distance_factor + elif forward_similarity > 0.0: + # Cameras looking in somewhat similar directions + distance_factor = 1.0 - normalized_distance * 0.8 + overlap = 0.05 + 0.15 * distance_factor * forward_similarity + else: + # Cameras looking away from each other - very low overlap + # But still give some score based on distance (closer = slightly better) + overlap = max(0.0, 0.01 - normalized_distance * 0.01) + + # Apply distance penalty for cameras exceeding max_distance (soft penalty) + if distance_exceeds_max: + # Reduce score by distance penalty, but don't make it zero + distance_penalty = min(0.5, (distance - max_distance) / max_distance * 0.3) + overlap = overlap * (1.0 - distance_penalty) + + return np.clip(overlap, 0.0, 1.0) + + +# Keep the simple version as fallback +def compute_fov_overlap_simple( + pose1: Dict, + pose2: Dict, + fov_degrees: float = 52.67, + max_distance: float = 50.0 +) -> float: + """ + Simplified FOV overlap computation (fallback). + Use compute_fov_overlap_3d for more accurate results. + """ + return compute_fov_overlap_3d(pose1, pose2, fov_degrees, max_distance) + + +class FOVMemoryRetriever: + """ + FOV-based Memory Retriever for Context-as-Memory. + + Retrieves relevant historical frames based on FOV overlap with current frame. + """ + + def __init__( + self, + dataset_base_path: str, + fov_degrees: float = 52.67, + max_distance: float = 50.0, + use_precomputed_overlaps: bool = True + ): + """ + Initialize FOV Memory Retriever. + + Args: + dataset_base_path: Base path to Context-as-Memory dataset + fov_degrees: Field of view in degrees (default from paper: 52.67) + max_distance: Maximum distance to consider for overlap (meters) + use_precomputed_overlaps: Whether to use precomputed overlap_labels if available + """ + self.dataset_base_path = dataset_base_path + self.fov_degrees = fov_degrees + self.max_distance = max_distance + self.use_precomputed_overlaps = use_precomputed_overlaps + + self.jsons_dir = os.path.join(dataset_base_path, 'jsons') + self.overlap_labels_dir = os.path.join(dataset_base_path, 'overlap_labels') + + # Cache for loaded poses + self._pose_cache: Dict[str, Dict] = {} + + def retrieve_frames( + self, + video_name: str, + current_frame_idx: int, + candidate_frame_indices: List[int], + top_k: int = 5, + include_last_frame: bool = True, + use_relative_poses: bool = False # Experiment 1_4_2: use RT relative conversion + ) -> List[int]: + """ + Retrieve top-k most relevant frames based on FOV overlap. + + According to Context-as-Memory, for temporal coherence, we should: + 1. Always include the last frame (current_frame_idx - 1) as short-term memory + 2. Retrieve top-(k-1) frames from history as long-term memory + + Args: + video_name: Name of the video + current_frame_idx: Index of current frame to generate + candidate_frame_indices: List of candidate frame indices to consider + top_k: Number of frames to retrieve + include_last_frame: Whether to force include the last frame (default: True, per Context-as-Memory) + use_relative_poses: Whether to use RT relative conversion (experiment 1_4_2, aligned with paper) + + Returns: + List of top-k frame indices sorted by relevance (last frame first if included) + """ + if not candidate_frame_indices: + return [] + + retrieved_frames = [] + + # Step 1: Force include last frame for short-term memory (Context-as-Memory requirement) + if include_last_frame and current_frame_idx > 0: + last_frame_idx = current_frame_idx - 1 + if last_frame_idx in candidate_frame_indices: + retrieved_frames.append(last_frame_idx) + # Remove from candidates to avoid duplication + candidate_frame_indices = [idx for idx in candidate_frame_indices if idx != last_frame_idx] + + # Calculate how many more frames we need + remaining_k = top_k - len(retrieved_frames) + if remaining_k <= 0: + return retrieved_frames[:top_k] + + # Step 2: Retrieve long-term memory frames using FOV overlap + # If precomputed overlaps are available, use them + if self.use_precomputed_overlaps: + overlap_frames = load_overlap_frames( + self.overlap_labels_dir, + video_name, + current_frame_idx + ) + + # Filter to only include candidate frames (and exclude already included last frame) + overlap_frames = [f for f in overlap_frames + if f in candidate_frame_indices and f not in retrieved_frames] + + if overlap_frames: + # Take top remaining_k frames + retrieved_frames.extend(overlap_frames[:remaining_k]) + return retrieved_frames[:top_k] + + # Step 3: Compute FOV overlap using camera poses (if precomputed not available) + current_pose = self._load_pose(video_name, current_frame_idx) + if current_pose is None: + # Fallback: return first k candidates + retrieved_frames.extend(candidate_frame_indices[:remaining_k]) + return retrieved_frames[:top_k] + + # Experiment 1_4_2: Convert to relative poses if enabled (aligned with Context-as-Memory) + if use_relative_poses: + # Convert current pose to RT format + ref_rt = pose_to_rt(current_pose) + if ref_rt is None: + # Fallback to absolute poses if conversion fails + use_relative_poses = False + + # Compute overlap scores for all candidates + overlap_scores = [] + for candidate_idx in candidate_frame_indices: + if candidate_idx in retrieved_frames: + continue # Skip already included frames + + candidate_pose = self._load_pose(video_name, candidate_idx) + if candidate_pose is None: + continue + + # Experiment 1_4_2: Use relative poses for FOV overlap computation + if use_relative_poses and ref_rt is not None: + # Convert candidate pose to RT format + candidate_rt = pose_to_rt(candidate_pose) + if candidate_rt is not None: + # Convert to relative coordinates + relative_rt_list = convert_rt_to_relative([candidate_rt], ref_rt) + if relative_rt_list: + # Convert back to pose format for FOV overlap computation + relative_pose = rt_to_pose(relative_rt_list[0]) + if relative_pose is not None: + # Use relative pose for overlap computation + # Reference pose in relative coordinates is identity (origin) + ref_relative_pose = {'position': [0, 0, 0], 'rotation': [0, 0, 0]} + score = compute_fov_overlap_3d( + ref_relative_pose, + relative_pose, + self.fov_degrees, + self.max_distance + ) + overlap_scores.append((candidate_idx, score)) + continue + + # Fallback: Use absolute poses (original method) + score = compute_fov_overlap_3d( + current_pose, + candidate_pose, + self.fov_degrees, + self.max_distance + ) + overlap_scores.append((candidate_idx, score)) + + # Sort by score (descending) and take top remaining_k + overlap_scores.sort(key=lambda x: x[1], reverse=True) + retrieved_frames.extend([idx for idx, _ in overlap_scores[:remaining_k]]) + + return retrieved_frames[:top_k] + + def _load_pose(self, video_name: str, frame_idx: int) -> Optional[Dict]: + """Load and cache camera pose.""" + cache_key = f"{video_name}_{frame_idx}" + if cache_key in self._pose_cache: + return self._pose_cache[cache_key] + + json_file = os.path.join(self.jsons_dir, f"{video_name}.json") + pose = load_camera_pose(json_file, frame_idx) + + if pose is not None: + self._pose_cache[cache_key] = pose + + return pose + + def clear_cache(self): + """Clear pose cache.""" + self._pose_cache.clear() + + +def create_fov_retriever(dataset_base_path: str) -> Optional[FOVMemoryRetriever]: + """ + Create FOV retriever if dataset has camera pose information. + + Args: + dataset_base_path: Base path to dataset + + Returns: + FOVMemoryRetriever instance or None if dataset doesn't support it + """ + jsons_dir = os.path.join(dataset_base_path, 'jsons') + if not os.path.exists(jsons_dir): + return None + + return FOVMemoryRetriever(dataset_base_path) + + +# Context/FOV retrieval integration helpers. +def _load_frame_png(frame_file: str) -> Optional[Image.Image]: + """Load a single frame from PNG file.""" + if os.path.exists(frame_file): + try: + return Image.open(frame_file).convert('RGB') + except Exception: + pass + return None + + + +def retrieve_simple_context_frames( + data: Dict, + dataset_base_path: str, + top_k: int = 4, # Number of overlap frames to retrieve. First Frame will be added automatically. + drop_overlap_probability: float = 0.1, # 10% probability to drop overlap frames (paper strategy) + use_rt_relative: bool = False, # Experiment 1_4_2: use RT relative conversion (aligned with Context-as-Memory) +) -> Tuple[List[Image.Image], List, List[int], int, str, str]: + """ + Retrieve context frames according to the Context-as-Memory paper [2506.03141]. + + Data Structure (Precomputed Retrieval Results): + - Each JSON file: overlap_labels/{video_name}/{frame_index}.json + - JSON structure: { + "frame_index": "0", # First Frame (short-term memory) + "overlapping_frames": ["2796", "2797", ..., "3839", ...] # Long-term memory (sample 4) + } + - frame_index and the next 80 frames constitute GT (target frames to generate) + - Iterating through all JSON files = one epoch + + Design principles: + 1. First Frame (frame_index from JSON) as Immediate Condition: + - The frame_index in JSON is the First Frame (short-term memory) + - Provides immediate visual and temporal starting point (Image-to-Video mode) + - Always included as context + - GT: frame_index and the next 80 frames (81 frames total) + + 2. Overlap Frames (from overlapping_frames in JSON) as Long-term Memory: + - Retrieved from precomputed overlap_labels JSON files + - Precomputed lists may be very long (e.g., [2796, 3839, 4183, ..., 6339]) + - Random uniform sampling: sample top_k (4) frames from overlapping_frames list + - Provides long-term consistency information + - Memory frames are unordered snapshots (no temporal sequence) + + 3. Context Composition: + - Order: [First Frame, Overlap Frame 1, Overlap Frame 2, Overlap Frame 3, Overlap Frame 4] + - Total: 1 First Frame + top_k Overlap Frames = top_k + 1 frames (e.g., 5 frames) + - Context frames are concatenated with target frames in temporal dimension + + 4. 10% Probability Drop Strategy: + - With 10% probability, drop all Overlap Frames, only use First Frame + - Simulates video generation starting stage (no historical memory) + - Forces model to generate reasonable videos without long-term memory assistance + + 5. Epoch Definition: + - One epoch = iterate through all JSON files in overlap_labels/{video_name}/ + - Each JSON file = one training sample + + 6. Positional Encoding Note: + - Memory frames should NOT use original absolute time positions + - They should be treated as unordered image collection or use memory ID encoding only + - First Frame should use explicit "Ref Frame" encoding + + Args: + data: Training data dict containing video frames and metadata + dataset_base_path: Base path to Context-as-Memory dataset + top_k: Number of overlap frames to retrieve (default: 4). First Frame will be added automatically. + drop_overlap_probability: Probability to drop overlap frames (default: 0.1 = 10%) + + Returns: + Tuple of: + (context_frames, context_actions, context_indices, current_frame_idx, video_name, source) + """ + video_frames = data.get("video", []) + # Get segment boundaries + start_frame = data.get("start_frame", 0) + end_frame = data.get("end_frame", None) + + # Use frame_idx if available, otherwise calculate from segment (middle of segment) + # This ensures we can find previous frames for context retrieval + if "frame_idx" in data: + current_frame_idx = data.get("frame_idx") + else: + # Calculate middle of segment as reference frame (same as training convention) + if end_frame is not None: + current_frame_idx = (start_frame + end_frame) // 2 + else: + # Fallback: use middle of video_frames if available + if len(video_frames) > 0: + current_frame_idx = len(video_frames) // 2 + else: + current_frame_idx = 0 + + # First Frame: current segment's first frame (start_frame) - ALWAYS included + first_frame_idx = start_frame + + video_name = data.get("video_name", "") + context_frames: List[Image.Image] = [] + context_actions: List = [] + context_indices: List[int] = [] + source = "none" + + # Get video frames from data + if not isinstance(video_frames, list): + video_frames = [] + + if not video_name: + # Try to infer from data + if "video_path" in data: + video_name = os.path.basename(data["video_path"]).replace(".mp4", "").replace(".avi", "") + elif "file_path" in data: + video_name = os.path.basename(data["file_path"]).replace(".mp4", "").replace(".avi", "") + + # Step 1: Load First Frame (current segment's first frame) - ALWAYS included + # According to JSON structure: frame_index in JSON is the First Frame (short-term memory) + # frame_index and the next 80 frames constitute GT (target frames to generate) + # First Frame provides immediate visual and temporal starting point (Image-to-Video mode) + frames_dir = os.path.join(dataset_base_path, 'frames', video_name) + first_frame_loaded = False + + # Experiment 1_4_2: Load camera poses once for all context frames (one JSON read) + json_file = os.path.join(dataset_base_path, "jsons", f"{video_name}.json") + poses_dict = load_poses_dict(json_file) + first_frame_pose_rt = None + first_frame_pose = poses_dict.get(str(first_frame_idx)) + if first_frame_pose is not None: + first_frame_pose_rt = pose_to_rt(first_frame_pose) + if first_frame_pose_rt is not None: + # Experiment 1_4_2: When use_rt_relative, first frame = reference frame = identity RT + # Target actions use ref=first_frame, so context first frame must also be identity + # to align with target's coordinate system (same frame = same RT representation) + if use_rt_relative: + # Identity RT: [t=0,0,0, R=eye(3)] = [0,0,0,1,0,0,0,1,0,0,0,1] + context_actions.append([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0]) + else: + context_actions.append(first_frame_pose_rt) + else: + context_actions.append([0.0] * 12) + else: + context_actions.append([0.0] * 12) + + if os.path.exists(frames_dir): + first_frame_file = os.path.join(frames_dir, f"{first_frame_idx:04d}.png") + if os.path.exists(first_frame_file): + try: + first_frame = Image.open(first_frame_file).convert('RGB') + context_frames.append(first_frame) + context_indices.append(first_frame_idx) + first_frame_loaded = True + except Exception as e: + pass # Frame loading failed, skip + + # Fallback: use first frame from video_frames if file not found + if not first_frame_loaded and len(video_frames) > 0: + if isinstance(video_frames[0], Image.Image): + context_frames.append(video_frames[0]) + context_indices.append(start_frame) + first_frame_loaded = True + elif isinstance(video_frames[0], str): + try: + first_frame = Image.open(video_frames[0]).convert('RGB') + context_frames.append(first_frame) + context_indices.append(start_frame) + first_frame_loaded = True + except Exception as e: + pass # Fallback frame loading failed, skip + + # Step 2: Retrieve Overlap Frames (long-term memory) - with 10% probability drop + # According to the data structure: + # - Each JSON file in overlap_labels/{video_name}/{frame_index}.json contains: + # - frame_index: First Frame (short-term memory) - this is the current segment's start + # - overlapping_frames: List of historical frames (long-term memory) - sample 4 from this list + # - frame_index and the next 80 frames constitute GT (target frames to generate) + # - Iterating through all JSON files = one epoch + + drop_overlap = random.random() < drop_overlap_probability + + if not drop_overlap: + # Use precomputed overlap_labels for FOV overlap-based selection + overlap_labels_dir = os.path.join(dataset_base_path, 'overlap_labels') + if os.path.exists(overlap_labels_dir): + # Load overlapping_frames from JSON file: overlap_labels/{video_name}/{current_frame_idx}.json + # The JSON structure: {"frame_index": "...", "overlapping_frames": ["2796", "2797", ...]} + overlapping_indices = load_overlap_frames( + overlap_labels_dir, + video_name, + current_frame_idx + ) + # Filter: exclude first_frame_idx, but allow frames with FOV overlap (may include future frames) + # FOV overlap indicates visual similarity, which is valuable for context memory + # even if the overlapping frame is from the future + overlapping_indices = [idx for idx in overlapping_indices + if idx != first_frame_idx] + + # If we have overlap frames, randomly sample top_k from them + # Note: overlap_labels are precomputed FOV overlap results, containing potentially + # very long lists of non-contiguous frame indices (e.g., [2796, 3839, 4183, ...]). + # We cannot use all frames due to memory constraints, so we use random uniform sampling. + # Sampling all JSON files once = one epoch + if overlapping_indices: + # Experiment 1_4_2: Use RT relative conversion for better geometric consistency + # Aligned with Context-as-Memory [2506.03141] - use relative camera poses for FOV overlap + if use_rt_relative: + ref_pose = poses_dict.get(str(current_frame_idx)) + + if ref_pose is not None: + # Convert reference pose to RT format + ref_rt = pose_to_rt(ref_pose) + + if ref_rt is not None: + # Score candidates using relative poses + candidate_scores = [] + for candidate_idx in overlapping_indices: + candidate_pose = poses_dict.get(str(candidate_idx)) + if candidate_pose is None: + continue + + # Convert to RT and compute relative pose + candidate_rt = pose_to_rt(candidate_pose) + if candidate_rt is not None: + relative_rt_list = convert_rt_to_relative([candidate_rt], ref_rt) + if relative_rt_list: + relative_pose = rt_to_pose(relative_rt_list[0]) + if relative_pose is not None: + # Compute FOV overlap using relative poses + ref_relative_pose = {'position': [0, 0, 0], 'rotation': [0, 0, 0]} + score = compute_fov_overlap_3d( + ref_relative_pose, + relative_pose, + fov_degrees=52.67, + max_distance=500.0 + ) + candidate_scores.append((candidate_idx, score)) + + # Sort by score and select top_k + if candidate_scores: + candidate_scores.sort(key=lambda x: x[1], reverse=True) + sampled_overlap_indices = [idx for idx, _ in candidate_scores[:top_k]] + else: + # Fallback to random sampling if RT conversion fails + num_overlap_frames = max(1, top_k) + sampled_overlap_indices = random.sample( + overlapping_indices, + min(len(overlapping_indices), num_overlap_frames) + ) + else: + # Fallback to random sampling if RT conversion fails + num_overlap_frames = max(1, top_k) + sampled_overlap_indices = random.sample( + overlapping_indices, + min(len(overlapping_indices), num_overlap_frames) + ) + else: + # Fallback to random sampling if reference pose not found + num_overlap_frames = max(1, top_k) + sampled_overlap_indices = random.sample( + overlapping_indices, + min(len(overlapping_indices), num_overlap_frames) + ) + else: + # Original strategy: Random Uniform Sampling (recommended for robustness) + # This allows the model to learn from diverse temporal spans of memory. + # The precomputed overlapping_frames list may contain hundreds or thousands of indices, + # but we only sample top_k (e.g., 4) frames due to memory constraints. + num_overlap_frames = max(1, top_k) + + # Random uniform sampling from the precomputed overlapping_frames list + # This treats memory frames as an unordered image collection + # IMPORTANT: Do NOT sort sampled_indices - they are unordered snapshots, not a temporal sequence + # Positional encoding should NOT use original absolute time positions for memory frames + sampled_overlap_indices = random.sample( + overlapping_indices, + min(len(overlapping_indices), num_overlap_frames) + ) + + # Load overlap frames (memory frames) - long-term memory + # Experiment 1_4_2: Use first_frame (start_frame) as reference for ALL context RTs + # Target actions use ref=first_frame; context must use same ref for trajectory alignment + ref_pose_for_rt = first_frame_pose_rt if use_rt_relative else None + + if os.path.exists(frames_dir): + to_load = [(idx, os.path.join(frames_dir, f"{idx:04d}.png")) for idx in sampled_overlap_indices[:top_k]] + frames_loaded = {} + with ThreadPoolExecutor(max_workers=max(1, min(5, len(to_load)))) as ex: + futures = {ex.submit(_load_frame_png, fp): idx for idx, fp in to_load} + for fut in as_completed(futures): + frame_idx = futures[fut] + frame = fut.result() + if frame is not None: + frames_loaded[frame_idx] = frame + for frame_idx in sampled_overlap_indices[:top_k]: + if frame_idx not in frames_loaded: + continue + frame = frames_loaded[frame_idx] + context_frames.append(frame) + context_indices.append(frame_idx) + pose = poses_dict.get(str(frame_idx)) + if pose is not None: + rt_pose = pose_to_rt(pose) + if rt_pose is not None: + if use_rt_relative and ref_pose_for_rt is not None: + relative_rt_list = convert_rt_to_relative([rt_pose], ref_pose_for_rt) + context_actions.append(relative_rt_list[0] if relative_rt_list else rt_pose) + else: + context_actions.append(rt_pose) + else: + context_actions.append([0.0] * 12) + else: + context_actions.append([0.0] * 12) + + source = "overlap_labels_random" + else: + # No overlap frames found, will use fallback below + source = "first_frame_only" + else: + # No overlap_labels directory, will use fallback below + source = "first_frame_only" + else: + # 10% probability: drop overlap frames, only use First Frame + # This simulates video generation starting stage (no historical memory) + source = "first_frame_only_dropped" + # Step 3: Fallback if we don't have enough overlap frames (and not dropped) + # Only fill if we haven't dropped overlap frames and need more frames + if source not in ["first_frame_only_dropped"] and len(context_frames) < top_k + 1: + # Try random fallback: use random previous frames before current_frame_idx + max_prev_frame = max(1, current_frame_idx - 1) + if max_prev_frame > 1: + # Exclude first_frame_idx from random sampling + candidate_indices = [idx for idx in range(max_prev_frame) + if idx != first_frame_idx and idx < current_frame_idx] + if candidate_indices: + num_needed = top_k + 1 - len(context_frames) + num_random_frames = min(len(candidate_indices), num_needed) + if num_random_frames > 0: + sampled_indices = random.sample(candidate_indices, num_random_frames) + + # Load random frames + if os.path.exists(frames_dir): + for frame_idx in sampled_indices: + frame_file = os.path.join(frames_dir, f"{frame_idx:04d}.png") + if os.path.exists(frame_file): + try: + frame = Image.open(frame_file).convert('RGB') + context_frames.append(frame) + context_indices.append(frame_idx) + pose = poses_dict.get(str(frame_idx)) + if pose is not None: + rt_pose = pose_to_rt(pose) + if rt_pose is not None: + # Convert to relative RT if enabled + if use_rt_relative and first_frame_pose_rt is not None: + relative_rt_list = convert_rt_to_relative([rt_pose], first_frame_pose_rt) + context_actions.append(relative_rt_list[0] if relative_rt_list else rt_pose) + else: + context_actions.append(rt_pose) + else: + context_actions.append([0.0] * 12) + else: + context_actions.append([0.0] * 12) + except Exception as e: + pass # Frame loading failed, skip + + if source == "first_frame_only": + source = "random_fallback" + + # Step 4: Final fallback - use additional frames from current segment if needed + # Target is top_k + 1 frames (1 First Frame + top_k Overlap/Random Frames) + target_total_frames = top_k + 1 + if len(context_frames) < target_total_frames and len(video_frames) > 0: + num_needed = target_total_frames - len(context_frames) + # Use additional frames from current segment (after first frame) + segment_start_idx = 1 if len(context_frames) > 0 else 0 + for i in range(segment_start_idx, min(segment_start_idx + num_needed, len(video_frames))): + frame_idx_seg = start_frame + i + if isinstance(video_frames[i], Image.Image): + context_frames.append(video_frames[i]) + context_indices.append(frame_idx_seg) + elif isinstance(video_frames[i], str): + # If it's a path, load it + try: + frame = Image.open(video_frames[i]).convert('RGB') + context_frames.append(frame) + context_indices.append(frame_idx_seg) + except: + pass + + pose = poses_dict.get(str(frame_idx_seg)) + if pose is not None: + rt_pose = pose_to_rt(pose) + if rt_pose is not None: + # Convert to relative RT if enabled + if use_rt_relative and first_frame_pose_rt is not None: + relative_rt_list = convert_rt_to_relative([rt_pose], first_frame_pose_rt) + context_actions.append(relative_rt_list[0] if relative_rt_list else rt_pose) + else: + context_actions.append(rt_pose) + else: + context_actions.append([0.0] * 12) + else: + context_actions.append([0.0] * 12) + + # Update source if we used segment frames + if len(context_frames) >= num_needed and source in ["first_frame_only", "none"]: + source = "segment_fallback" + + # Step 5: Ensure we return exactly top_k + 1 frames + # If we have some frames but not enough, pad by repeating the last frame + if len(context_frames) < target_total_frames: + if context_frames: + last_frame = context_frames[-1] + last_idx = context_indices[-1] if context_indices else first_frame_idx + last_action = context_actions[-1] if context_actions else [0.0] * 12 + while len(context_frames) < target_total_frames: + context_frames.append(last_frame) + context_indices.append(last_idx) + context_actions.append(last_action) # Always pad with pose data + + # Limit to top_k + 1 (in case we have more) + context_frames = context_frames[:target_total_frames] + if context_indices: + context_indices = context_indices[:target_total_frames] + if context_actions: + context_actions = context_actions[:target_total_frames] + + # Ensure context_actions length matches context_frames (pad with zeros if needed) + # Always ensure context_actions are provided (not just when use_rt_relative=True) + while len(context_actions) < len(context_frames): + context_actions.append([0.0] * 12) + + return context_frames, context_actions, context_indices, current_frame_idx, video_name, source + + +def retrieve_fov_context_frames( + data: Dict, + dataset_base_path: str, + fov_retriever=None, + top_k: int = 4, # Number of overlap frames to retrieve. First Frame will be added automatically. + use_precomputed_overlaps: bool = True, + use_rt_relative: bool = False, # Experiment 1_4_2: Use RT relative conversion (aligned with Context-as-Memory) + strict_overlap_labels: bool = False, + allow_realtime_fallback: bool = True, + allow_segment_fallback: bool = True, + drop_overlap_probability: float = 0.1, # 10% probability to drop overlap frames (paper strategy) +): + """ + Backward-compatible wrapper. + We use FOV overlap scoring to select top-k overlap frames (as in Context-as-Memory). + + According to Context-as-Memory [2506.03141]: + - First Frame (current segment's first frame) is always included as immediate condition + - Overlap Frames are retrieved as long-term memory + - With 10% probability, drop overlap frames to simulate starting stage + + Experiment 1_4_2: Uses RT relative conversion for better geometric consistency. + """ + context_frames, context_actions, context_indices, cur_idx, video_name, source = retrieve_simple_context_frames( + data=data, + dataset_base_path=dataset_base_path, + top_k=top_k, # top_k is number of overlap frames (4), First Frame will be added automatically (total: 5) + use_rt_relative=use_rt_relative, # Experiment 1_4_2: RT relative conversion + drop_overlap_probability=drop_overlap_probability, # 10% probability to drop overlap frames + ) + # Check if we have top_k + 1 frames (1 First Frame + top_k Overlap Frames) + target_total_frames = top_k + 1 + if strict_overlap_labels and len(context_frames) < target_total_frames: + return [], [], [], cur_idx, video_name, "overlap_labels_insufficient" + return context_frames, context_actions, context_indices, cur_idx, video_name, source + + +def save_sampling_jsonl( + output_path: str, + video_name: str, + frame_index: int, + context_indices: List[int], + prompt: Optional[str] = None, + start_frame: Optional[int] = None, + end_frame: Optional[int] = None, + source: Optional[str] = None, + append: bool = True, +) -> None: + """ + Save context sampling result to JSONL file for eval consistency. + + Format: + { + "video_name": "AncientTempleEnv_0", + "frame_index": 0, # First Frame (short-term memory) + "context_indices": [0, 2796, 3839, 4183, 6339], # First Frame + 4 Overlap Frames + "prompt": "...", # Optional + "start_frame": 0, # Optional: GT segment start + "end_frame": 80, # Optional: GT segment end + "source": "overlap_labels_random" # Optional: sampling source + } + + Args: + output_path: Path to JSONL file + video_name: Video name + frame_index: First Frame index (from JSON file) + context_indices: List of context frame indices [first_frame, overlap1, overlap2, ...] + prompt: Optional prompt text + start_frame: Optional GT segment start frame + end_frame: Optional GT segment end frame + source: Optional sampling source + append: Whether to append to existing file (default: True) + """ + os.makedirs(os.path.dirname(output_path), exist_ok=True) + + item = { + "video_name": video_name, + "frame_index": frame_index, + "context_indices": context_indices, + } + + if prompt is not None: + item["prompt"] = prompt + if start_frame is not None: + item["start_frame"] = start_frame + if end_frame is not None: + item["end_frame"] = end_frame + if source is not None: + item["source"] = source + + mode = "a" if append else "w" + with open(output_path, mode, encoding="utf-8") as f: + f.write(json.dumps(item, ensure_ascii=False) + "\n") + + +def load_sampling_jsonl(jsonl_path: str) -> List[Dict]: + """ + Load context sampling results from JSONL file. + + Args: + jsonl_path: Path to JSONL file + + Returns: + List of sampling items, each containing: + { + "video_name": str, + "frame_index": int, + "context_indices": List[int], + "prompt": Optional[str], + "start_frame": Optional[int], + "end_frame": Optional[int], + "source": Optional[str] + } + """ + if not os.path.exists(jsonl_path): + return [] + + items = [] + with open(jsonl_path, "r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + try: + item = json.loads(line) + items.append(item) + except json.JSONDecodeError as e: + print(f"Warning: Failed to parse JSONL line: {e}") + continue + + return items + + +def setup_fov_retriever_for_training( + dataset_base_path: str, + enable_fov_retrieval: bool = True, +) -> Optional[object]: + """ + Setup FOV retriever for training (simplified version). + Context frames are selected by FOV overlap scoring from precomputed `overlap_labels`. + """ + return None + + +_WARN_ONCE_KEYS = set() + + +def _warn_once(key: str, msg: str) -> None: + if key in _WARN_ONCE_KEYS: + return + _WARN_ONCE_KEYS.add(key) + print(f"[context_retrieval] WARN: {msg}", file=sys.stderr, flush=True) + + +def _load_latent(latent_dir: str, video_name: str, frame_idx: int): + """Load a single-frame latent. Expects latent_dir/video_name/{frame_idx:04d}.pt or .pt with key 'latent' or raw tensor.""" + base = os.path.join(latent_dir, video_name) + for fmt in (f"{frame_idx:04d}.pt", f"{frame_idx}.pt"): + path = os.path.join(base, fmt) + if os.path.isfile(path): + try: + try: + x = torch.load(path, map_location="cpu", weights_only=True) + except TypeError: + x = torch.load(path, map_location="cpu") + if isinstance(x, dict) and "latent" in x: + z = x["latent"] + else: + z = x + if hasattr(z, "shape"): + # (C, 1, H, W) or (C, H, W) -> flatten for similarity + return z.flatten() + return None + except Exception: + pass + return None + + +def latent_sim_rank( + video_name: str, + first_frame_idx: int, + overlapping_indices: List[int], + dataset_base_path: str, + top_k: int, + latent_dir: Optional[str] = None, + use_cosine: bool = True, +) -> List[int]: + """ + Rank overlapping frame indices by latent similarity to the first (reference) frame. + If latent_dir is None or latents are missing, falls back to random sample. + Expects per-frame latents under latent_dir/video_name/{frame_idx:04d}.pt (or {frame_idx}.pt). + """ + if not overlapping_indices or top_k <= 0: + return [] + if latent_dir is None or not os.path.isdir(latent_dir): + return random.sample(overlapping_indices, min(top_k, len(overlapping_indices))) + + ref = _load_latent(latent_dir, video_name, first_frame_idx) + if ref is None: + return random.sample(overlapping_indices, min(top_k, len(overlapping_indices))) + + ref = ref.float().unsqueeze(0) + scores = [] + for idx in overlapping_indices: + cand = _load_latent(latent_dir, video_name, idx) + if cand is None: + continue + cand = cand.float().unsqueeze(0) + if use_cosine: + sim = torch.nn.functional.cosine_similarity(ref, cand, dim=1).item() + else: + sim = -((ref - cand) ** 2).sum().item() + scores.append((idx, sim)) + if not scores: + return random.sample(overlapping_indices, min(top_k, len(overlapping_indices))) + scores.sort(key=lambda x: x[1], reverse=True) + return [idx for idx, _ in scores[:top_k]] + + +def retrieve_context_frames_advanced( + data: Dict, + dataset_base_path: str, + top_k: int = 4, + drop_overlap_probability: float = 0.1, + use_rt_relative: bool = False, + retrieval_method: str = "fov", + latent_retrieval_dir: Optional[str] = None, + strict_overlap_labels: bool = False, +) -> Tuple[List, List, List[int], int, str, str]: + """ + Retrieve context frames with pluggable retrieval method. + Interface matches retrieve_fov_context_frames return: + (context_frames, context_actions, context_indices, cur_idx, video_name, source). + + retrieval_method: + - "fov": use existing FOV/overlap_labels logic (random or RT-scored from overlap). + - "latent_sim": rank overlap candidates by latent similarity to first frame (requires latent_retrieval_dir). + When latent_sim is used but latent_retrieval_dir is missing or latents absent, falls back to FOV behavior. + """ + if retrieval_method == "latent_sim" and not latent_retrieval_dir: + _warn_once( + "latent_sim_missing_dir", + "retrieval_method=latent_sim but latent_retrieval_dir is not set; fallback to FOV retrieval.", + ) + if retrieval_method == "fov" or (retrieval_method == "latent_sim" and not latent_retrieval_dir): + return retrieve_simple_context_frames( + data=data, + dataset_base_path=dataset_base_path, + top_k=top_k, + drop_overlap_probability=drop_overlap_probability, + use_rt_relative=use_rt_relative, + ) + if retrieval_method == "latent_sim" and latent_retrieval_dir and not os.path.isdir(latent_retrieval_dir): + _warn_once( + "latent_sim_bad_dir", + f"latent_retrieval_dir not found: {latent_retrieval_dir}; latent_sim will degrade to random-overlap selection.", + ) + + # latent_sim: we need to inject a custom ranking into the flow. We do a minimal duplicate of the + # overlap selection step then reuse the rest via a wrapper around retrieve_simple_context_frames + # by passing a custom rank function. Since retrieve_simple_context_frames doesn't support that yet, + # we implement a full path here that mirrors it but uses latent_sim_rank for overlap selection. + + video_frames = data.get("video", []) + start_frame = data.get("start_frame", 0) + end_frame = data.get("end_frame", None) + if "frame_idx" in data: + current_frame_idx = data.get("frame_idx") + else: + current_frame_idx = (start_frame + end_frame) // 2 if end_frame is not None else (len(video_frames) // 2 if video_frames else 0) + first_frame_idx = start_frame + video_name = data.get("video_name", "") + if not video_name and "video_path" in data: + video_name = os.path.basename(data["video_path"]).replace(".mp4", "").replace(".avi", "") + elif not video_name and "file_path" in data: + video_name = os.path.basename(data["file_path"]).replace(".mp4", "").replace(".avi", "") + + frames_dir = os.path.join(dataset_base_path, "frames", video_name) + json_file = os.path.join(dataset_base_path, "jsons", f"{video_name}.json") + poses_dict = load_poses_dict(json_file) if os.path.isfile(json_file) else {} + first_frame_pose = poses_dict.get(str(first_frame_idx)) + first_frame_pose_rt = pose_to_rt(first_frame_pose) if first_frame_pose is not None and pose_to_rt else None + + context_frames: List[Image.Image] = [] + context_actions: List = [] + context_indices: List[int] = [] + source = "none" + + def _append_pose(frame_idx: int): + pose = poses_dict.get(str(frame_idx)) + if pose is not None and pose_to_rt and use_rt_relative and first_frame_pose_rt is not None: + rt = pose_to_rt(pose) + if rt is not None and convert_rt_to_relative: + rel = convert_rt_to_relative([rt], first_frame_pose_rt) + context_actions.append(rel[0] if rel else [0.0] * 12) + else: + context_actions.append([0.0] * 12) + elif pose is not None and pose_to_rt: + rt = pose_to_rt(pose) + context_actions.append(rt if rt is not None else [0.0] * 12) + else: + context_actions.append([0.0] * 12) + + if first_frame_pose_rt is not None and use_rt_relative: + context_actions.append([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0]) + else: + context_actions.append(first_frame_pose_rt if first_frame_pose_rt is not None else [0.0] * 12) + + if os.path.isdir(frames_dir): + first_path = os.path.join(frames_dir, f"{first_frame_idx:04d}.png") + if os.path.isfile(first_path): + try: + context_frames.append(Image.open(first_path).convert("RGB")) + context_indices.append(first_frame_idx) + except Exception: + pass + if not context_frames and video_frames: + if isinstance(video_frames[0], Image.Image): + context_frames.append(video_frames[0]) + context_indices.append(start_frame) + elif isinstance(video_frames[0], str) and os.path.isfile(video_frames[0]): + try: + context_frames.append(Image.open(video_frames[0]).convert("RGB")) + context_indices.append(start_frame) + except Exception: + pass + + drop_overlap = random.random() < drop_overlap_probability + if drop_overlap: + target_total = top_k + 1 + if strict_overlap_labels and len(context_frames) < target_total: + return [], [], [], current_frame_idx, video_name, "overlap_labels_insufficient" + return context_frames, context_actions, context_indices, current_frame_idx, video_name, "first_frame_only_dropped" + + overlap_labels_dir = os.path.join(dataset_base_path, "overlap_labels") + overlapping_indices = [] + if os.path.isdir(overlap_labels_dir): + overlapping_indices = load_overlap_frames(overlap_labels_dir, video_name, current_frame_idx) + overlapping_indices = [i for i in overlapping_indices if i != first_frame_idx] + + if not overlapping_indices: + source = "first_frame_only" + if strict_overlap_labels and len(context_frames) < top_k + 1: + return [], [], [], current_frame_idx, video_name, "overlap_labels_insufficient" + return context_frames, context_actions, context_indices, current_frame_idx, video_name, source + + sampled_overlap_indices = latent_sim_rank( + video_name, first_frame_idx, overlapping_indices, dataset_base_path, top_k, + latent_dir=latent_retrieval_dir, use_cosine=True, + ) + + def _load_frame(path: str): + if os.path.isfile(path): + try: + return Image.open(path).convert("RGB") + except Exception: + pass + return None + + to_load = [(idx, os.path.join(frames_dir, f"{idx:04d}.png")) for idx in sampled_overlap_indices[:top_k]] + with ThreadPoolExecutor(max_workers=max(1, min(5, len(to_load)))) as ex: + futures = {ex.submit(_load_frame, path): idx for idx, path in to_load} + for fut in futures: + idx = futures[fut] + frame = fut.result() + if frame is not None: + context_frames.append(frame) + context_indices.append(idx) + _append_pose(idx) + + source = "latent_sim" + if strict_overlap_labels and len(context_frames) < top_k + 1: + return [], [], [], current_frame_idx, video_name, "overlap_labels_insufficient" + return context_frames, context_actions, context_indices, current_frame_idx, video_name, source diff --git a/code/src/model_training/multichunk_sample_utils.py b/code/src/model_training/multichunk_sample_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..eb0f6da56f17055003b1e9ed6d4ac60fdd31ef2e --- /dev/null +++ b/code/src/model_training/multichunk_sample_utils.py @@ -0,0 +1,877 @@ +""" +Shared multichunk sampling for training monitor and replay scripts. + +Two-chunk path matches run_replay_loop_two_chunk: chunk1 with 1-frame context, +chunk2 with context_frames_for_next_chunk from chunk1 output. +No hidden state is carried across chunks (each pipe() is independent diffusion), only PIL frames + latents. +""" +from __future__ import annotations + +import argparse +import json +import os +import random +import traceback +from typing import Any, Dict, List, Optional, Sequence, Tuple + +import numpy as np +import torch +from PIL import Image +from diffsynth import save_video +from diffsynth.pipelines.wan_video_new import ModelConfig, WanVideoPipeline +from src.model_training.transformers_compat import patch_transformers_hybrid_cache + +patch_transformers_hybrid_cache() + +from diffsynth.trainers.utils import VideoDataset +from safetensors.torch import load_file as safe_load_file + +from src.model_training.fov_retrieval import load_camera_poses_batch +from src.model_training.fov_retrieval import convert_rt_to_relative, pose_to_rt + + +FrameType = Any + + +def context_frames_for_next_chunk(frames_list: Sequence[FrameType], K: int) -> List[FrameType]: + """Select K context frames from a finished chunk for the next chunk (replay-style). + + Order is [last_frame, ...]: last frame first (adjacent to target), then K-1 uniformly + spaced frames from indices [0 .. n-2]. + + - K==1: [last] + - K>1: [last] + (K-1) uniform samples from [0, n-2] + """ + n = len(frames_list) + if n <= 0 or K <= 0: + return [] + if K == 1: + return [frames_list[-1]] + n_ctx = min(K, n) + if n_ctx == 1: + return [frames_list[-1]] + last = frames_list[-1] + num_rest = n_ctx - 1 + if num_rest <= 0: + return [last] + if num_rest == 1: + return [last, frames_list[0]] + indices = [int(round(i * (n - 2) / (num_rest - 1))) for i in range(num_rest)] + rest = [frames_list[i] for i in indices] + return [last] + rest + + +def replay_context_global_indices(n_frames: int, K: int) -> List[int]: + """Indices into frames_list matching context_frames_for_next_chunk order (for tests/debug).""" + if n_frames <= 0 or K <= 0: + return [] + if K == 1: + return [n_frames - 1] + n_ctx = min(K, n_frames) + if n_ctx == 1: + return [n_frames - 1] + num_rest = n_ctx - 1 + if num_rest == 1: + return [n_frames - 1, 0] + indices = [int(round(i * (n_frames - 2) / (num_rest - 1))) for i in range(num_rest)] + return [n_frames - 1] + indices + + +def replay_context_from_generated_frames( + frames_list: Sequence[FrameType], + n_ctx: int, +) -> List[FrameType]: + """Single replay-style context selection entrypoint used by callsites. + + Keep legacy semantics: + - n_ctx > 0: replay sampling rule (last + uniform historical) + - n_ctx <= 0: fallback to last frame only + """ + n_ctx = int(n_ctx) + if n_ctx > 0: + return context_frames_for_next_chunk(frames_list, n_ctx) + return [frames_list[-1]] + + +def prev_chunk_tail_global_indices(start_frame: int, N: int, *, nearest_first: bool = False) -> Optional[List[int]]: + """Strict consecutive globals with configurable order. + + - nearest_first=False: [start_frame - N, ..., start_frame - 1] (oldest -> newest) + - nearest_first=True: [start_frame - 1, ..., start_frame - N] (newest -> oldest) + None if start_frame < N. + """ + if N <= 0: + return [] + if start_frame < N: + return None + if nearest_first: + return list(range(int(start_frame) - 1, int(start_frame) - N - 1, -1)) + return list(range(int(start_frame) - N, int(start_frame))) + + +def load_prev_chunk_tail_from_disk( + dataset_base_path: str, + video_name: str, + start_frame: int, + N: int, + *, + nearest_first: bool = False, +) -> Tuple[Optional[List[Any]], Optional[List[int]]]: + """Load N frames before start_frame in configured order.""" + + idxs = prev_chunk_tail_global_indices(int(start_frame), int(N), nearest_first=nearest_first) + if idxs is None: + return None, None + if not idxs: + return [], [] + vn = str(video_name) + if vn.endswith((".mp4", ".avi")): + vn = os.path.splitext(vn)[0] + frames_root = os.path.join(dataset_base_path, "frames", vn) + out: List[Any] = [] + for idx in idxs: + path = os.path.join(frames_root, f"{int(idx):04d}.png") + if not os.path.isfile(path): + return None, None + try: + out.append(Image.open(path).convert("RGB")) + except Exception: + return None, None + return out, idxs + + +def synthetic_replay_context_from_segment( + video_frames: Sequence[FrameType], + chunk_frames: int, + K: int, +) -> Optional[List[FrameType]]: + """Use first `chunk_frames` of video_frames as virtual chunk1; context for 'chunk2' via replay rule. + + Requires len(video_frames) >= chunk_frames. Returns None otherwise. + """ + if len(video_frames) < chunk_frames or K <= 0: + return None + chunk1 = list(video_frames[:chunk_frames]) + return context_frames_for_next_chunk(chunk1, K) + + +def replay_context_actions_from_segment_actions( + actions: Sequence[Sequence[float]], + n_frames: int, + K: int, +) -> Optional[List[List[float]]]: + """Align RT/action rows with context_frames_for_next_chunk order (same indices as replay_context_global_indices).""" + idxs = replay_context_global_indices(int(n_frames), int(K)) + if not idxs: + return [] + need_max = max(idxs) + if need_max >= len(actions): + return None + return [list(actions[i]) for i in idxs] + + +def load_prev_chunk_tail_rt_actions( + dataset_base_path: str, + video_name: str, + start_frame: int, + N: int, + *, + use_rt_relative: bool = True, + nearest_first: bool = False, +) -> Tuple[Optional[List[List[float]]], Optional[List[int]]]: + """Load RT poses in configured order, relative to first context frame.""" + idxs = prev_chunk_tail_global_indices(int(start_frame), int(N), nearest_first=nearest_first) + if idxs is None: + return None, None + if not idxs: + return [], [] + + vn = str(video_name) + if vn.endswith((".mp4", ".avi")): + vn = os.path.splitext(vn)[0] + json_file = os.path.join(dataset_base_path, "jsons", f"{vn}.json") + if not os.path.isfile(json_file): + return None, None + poses = load_camera_poses_batch(json_file, idxs) + rt_list = [pose_to_rt(p) if p else None for p in poses] + if not rt_list or any(r is None for r in rt_list): + return None, None + ref_rt = rt_list[0] + if use_rt_relative: + out = convert_rt_to_relative(rt_list, ref_rt) + else: + out = [list(r) for r in rt_list] + return out, idxs + +def encode_context_frames(pipe, pil_list, device, dtype=torch.bfloat16, per_frame: bool = False): + """Encode context frames to latents aligned with training behavior. + + per_frame=False: encode the whole clip once (default training path, temporal downsample). + per_frame=True: encode each frame separately and concat on latent time. + """ + if not pil_list: + return None + if not per_frame: + context_video = pipe.preprocess_video(pil_list).to(device=device) + if context_video.dim() == 5: + context_video = context_video.squeeze(0) + context_latents = pipe.vae.encode([context_video], device=pipe.device, tiled=False, tile_size=None, tile_stride=None) + return context_latents.to(dtype=dtype, device=device) + + encoded = [] + for pil in pil_list: + frame_video = pipe.preprocess_video([pil]).to(device=device) + frame_sq = frame_video.squeeze(0) if frame_video.dim() == 5 else frame_video + if frame_sq.dim() == 3: + frame_sq = frame_sq.unsqueeze(0) + lat_one = pipe.vae.encode([frame_sq], device=pipe.device, tiled=False, tile_size=None, tile_stride=None) + encoded.append(lat_one) + context_latents = torch.cat(encoded, dim=2).to(dtype=dtype, device=device) + return context_latents + + +def _frame_to_pil(f, tw, th): + if hasattr(f, "convert") and hasattr(f, "resize"): + return f.convert("RGB").resize((tw, th)) + if isinstance(f, np.ndarray): + if f.dtype != np.uint8: + f = (f * 255).astype(np.uint8) if f.max() <= 1.0 else f.astype(np.uint8) + return Image.fromarray(f).convert("RGB").resize((tw, th)) + if isinstance(f, torch.Tensor): + fn = f.cpu().numpy() + if len(fn.shape) == 3 and fn.shape[0] == 3: + fn = fn.transpose(1, 2, 0) + fn = (fn * 255).clip(0, 255).astype(np.uint8) if fn.max() <= 1.0 else fn.clip(0, 255).astype(np.uint8) + return Image.fromarray(fn).convert("RGB").resize((tw, th)) + return f + + +def run_one_chunk( + pipe, + prompt: str, + use_negative_prompt: str, + action_path: Optional[str] = None, + *, + cam_pose_actions=None, + context_latents=None, + num_context_frames: int = 1, + context_actions_t=None, + chunk_frames: int = 81, + h: int = 352, + w: int = 640, + seed: int = 0, + sigma_shift: float = 5.0, + num_inference_steps: int = 50, + cfg_scale: float = 5.0, + inference_noise_level: float = 0.0, + omit_context_actions: bool = False, # kept for backward compat, no longer used + context_position: str = "suffix", + log_prefix: str = "[multichunk]", +) -> List[Any]: + """Single chunk generation with explicit context position. VWM-aligned action injection.""" + device = pipe.device + kwargs_common = dict( + prompt=prompt, + negative_prompt=use_negative_prompt, + height=h, + width=w, + num_frames=chunk_frames, + num_inference_steps=num_inference_steps, + seed=seed, + cfg_scale=cfg_scale, + sigma_shift=sigma_shift, + denoising_strength=1.0, + ) + if action_path is not None: + kwargs_common["action_path"] = action_path + elif cam_pose_actions is not None: + kwargs_common["cam_pose_actions"] = cam_pose_actions + + if context_latents is not None: + pipe_kw = dict( + **kwargs_common, + enable_context_memory=True, + context_latents=context_latents, + num_context_frames=num_context_frames, + context_position=context_position, + cfg_target_only=True, + inference_noise_level=inference_noise_level, + ) + if context_actions_t is not None: + pipe_kw["context_actions"] = context_actions_t + with torch.no_grad(): + vid = pipe(**pipe_kw) + else: + with torch.no_grad(): + vid = pipe(**kwargs_common, enable_context_memory=False) + return vid if isinstance(vid, list) else [vid] + + +def _load_actions_tensor_from_json( + action_path: Optional[str], + *, + device: torch.device, + dtype: torch.dtype = torch.float32, +) -> Optional[torch.Tensor]: + if not action_path or not os.path.exists(action_path): + return None + try: + with open(action_path, "r", encoding="utf-8") as f: + data = json.load(f) + seq = data.get("actions", data) + items = sorted( + ((int(k), v) for k, v in seq.items() if str(k).isdigit()), + key=lambda x: x[0], + ) + if not items: + return None + rows = [] + for _, v in items: + if isinstance(v, (list, tuple)) and len(v) >= 12: + rows.append([float(x) for x in v[:12]]) + if not rows: + return None + return torch.tensor(rows, device=device, dtype=dtype) + except Exception: + return None + + +def _tail_context_actions( + src_actions: Optional[torch.Tensor], + num_ctx: int, + *, + device: torch.device, + dtype: torch.dtype = torch.float32, + nearest_first: bool = False, +) -> Optional[torch.Tensor]: + if num_ctx <= 0: + return None + if src_actions is None or src_actions.numel() == 0: + return torch.zeros(num_ctx, 12, device=device, dtype=dtype) + if src_actions.dim() == 3: + src_actions = src_actions[0] + if src_actions.shape[0] >= num_ctx: + out = src_actions[-num_ctx:] + if nearest_first: + out = torch.flip(out, dims=[0]) + return out.to(device=device, dtype=dtype) + pad_n = num_ctx - src_actions.shape[0] + pad = src_actions[-1:, :].expand(pad_n, src_actions.shape[1]) + out = torch.cat([src_actions, pad], dim=0) + if nearest_first: + out = torch.flip(out, dims=[0]) + return out.to(device=device, dtype=dtype) + + +def sync_pipe_memory_from_training_module(pipe, unwrapped_model: Any) -> Dict[str, Any]: + """Copy memory-related flags from WanTrainingModule.pipe onto pipe (defensive if pipe handle diverges).""" + log: Dict[str, Any] = {} + p = pipe + m = unwrapped_model + src = getattr(m, "pipe", None) or p + + def _g(attr, default=None): + v = getattr(src, attr, None) + if v is None: + v = getattr(p, attr, None) + if v is None: + v = getattr(m, attr, default) + return v + + p.use_framepack_memory = bool(_g("use_framepack_memory", False)) + p.context_temporal_decay = float(_g("context_temporal_decay", 1.0) or 1.0) + p.context_attention_weight = float(_g("context_attention_weight", 1.0) or 1.0) + p.use_framepack_length_compress = bool(_g("use_framepack_length_compress", False)) + p.framepack_ratio = int(_g("framepack_ratio", 1) or 1) + p.framepack_length_strategy = str(_g("framepack_length_strategy", "distance_merge") or "distance_merge") + p.framepack_recent_keep_ratio = float(_g("framepack_recent_keep_ratio", 0.5) or 0.5) + p.framepack_multiscale_w2 = float(_g("framepack_multiscale_w2", 0.25) or 0.25) + p.framepack_multiscale_w4 = float(_g("framepack_multiscale_w4", 0.15) or 0.15) + p.use_spatial_memory = bool(_g("use_spatial_memory", False)) + p.spatial_memory_tokens = int(_g("spatial_memory_tokens", 64) or 64) + p.use_spatial_memory_legacy = bool(_g("use_spatial_memory_legacy", False)) + p.spatial_memory_inject_mode = str(_g("spatial_memory_inject_mode", "concat_text") or "concat_text") + sm = getattr(m, "spatial_memory_module", None) or getattr(src, "spatial_memory_module", None) or getattr(p, "spatial_memory_module", None) + p.spatial_memory_module = sm + srm = getattr(m, "spatial_memory_readout_module", None) or getattr(src, "spatial_memory_readout_module", None) or getattr(p, "spatial_memory_readout_module", None) + p.spatial_memory_readout_module = srm + dit = getattr(p, "dit", None) + bl0 = dit.blocks[0] if dit is not None and hasattr(dit, "blocks") and len(dit.blocks) > 0 else None + log.update( + { + "use_framepack_memory": p.use_framepack_memory, + "use_framepack_length_compress": p.use_framepack_length_compress, + "framepack_ratio": p.framepack_ratio, + "framepack_length_strategy": p.framepack_length_strategy, + "use_spatial_memory": p.use_spatial_memory, + "use_spatial_memory_legacy": p.use_spatial_memory_legacy, + "spatial_memory_inject_mode": p.spatial_memory_inject_mode, + "spatial_module": sm is not None, + "spatial_readout_module": srm is not None, + "dit_block0_use_block_wise_ssm": bool(getattr(bl0, "use_block_wise_ssm", False)), + "dit_block0_use_videossm_hybrid": bool(getattr(bl0, "use_videossm_hybrid", False)), + } + ) + return log + + +def run_two_chunk_memory_monitor( + pipe, + *, + prompt: str, + negative_prompt: str, + action_path: Optional[str], + chunk0_action_path: Optional[str] = None, + chunk1_action_path: Optional[str] = None, + first_frame_pil, + context_memory_frames: int, + chunk_frames: int = 81, + h: int = 352, + w: int = 640, + seed: int = 42, + sigma_shift: float = 5.0, + num_inference_steps: int = 50, + cfg_scale: float = 5.0, + inference_noise_level: float = 0.0, + omit_context_actions: bool = False, + context_source: str = "replay", + context_position: str = "suffix", + context_per_frame_vae: bool = False, + device=None, + dtype=torch.bfloat16, + log_prefix: str = "[two_chunk_mem]", +) -> Tuple[List[Any], List[Any], Dict[str, Any]]: + """ + Chunk1: 1-frame context. Chunk2 context follows context_source: + - replay: context_frames_for_next_chunk + - prev_chunk_tail: strict tail frames (nearest-first) + + Returns (frames_ch0, frames_ch1, meta). chunk0 defaults left_45 and chunk1 defaults right_45 when provided by caller. + """ + device = device or pipe.device + context_source = (context_source or "replay").strip().lower() + if context_source not in ("replay", "prev_chunk_tail"): + context_source = "replay" + context_position = (context_position or "suffix").strip().lower() + if context_position not in ("prefix", "suffix"): + context_position = "suffix" + meta: Dict[str, Any] = { + "n_ctx": int(context_memory_frames), + "chunk_frames": chunk_frames, + "context_source": context_source, + "context_position": context_position, + "context_per_frame_vae": bool(context_per_frame_vae), + } + + ff = first_frame_pil + if isinstance(ff, Image.Image): + ff = ff.convert("RGB").resize((w, h), Image.Resampling.LANCZOS) + else: + ff = _frame_to_pil(ff, w, h) + + ctx_lat_0 = encode_context_frames(pipe, [ff], device, dtype=dtype, per_frame=bool(context_per_frame_vae)) + num_ctx0 = int(ctx_lat_0.shape[2]) if ctx_lat_0 is not None else 1 + meta["chunk0_num_context_latent"] = num_ctx0 + + use_omit_ch0 = omit_context_actions or (num_ctx0 <= 1) + act0 = chunk0_action_path or action_path + act1 = chunk1_action_path or action_path + src_actions0 = _load_actions_tensor_from_json(act0, device=device, dtype=torch.float32) + meta["chunk0_action_path"] = act0 + meta["chunk1_action_path"] = act1 + + frames_ch0 = run_one_chunk( + pipe, + prompt, + negative_prompt, + act0, + context_latents=ctx_lat_0, + num_context_frames=num_ctx0, + context_actions_t=None, + chunk_frames=chunk_frames, + h=h, + w=w, + seed=seed, + sigma_shift=sigma_shift, + num_inference_steps=num_inference_steps, + cfg_scale=cfg_scale, + inference_noise_level=inference_noise_level, + omit_context_actions=use_omit_ch0, + context_position=context_position, + log_prefix=log_prefix + " ch0", + ) + + pil_ch0 = [_frame_to_pil(f, w, h) for f in frames_ch0] + n_ctx = int(context_memory_frames) + if n_ctx <= 0: + n_ctx = 1 + if context_source == "prev_chunk_tail": + tail = pil_ch0[-n_ctx:] + prev_pil = list(reversed(tail)) if context_position == "suffix" else tail + else: + prev_pil = context_frames_for_next_chunk(pil_ch0, n_ctx) + meta["chunk1_context_count"] = len(prev_pil) + + ctx_lat_1 = encode_context_frames(pipe, prev_pil, device, dtype=dtype, per_frame=bool(context_per_frame_vae)) + num_ctx1 = int(ctx_lat_1.shape[2]) if ctx_lat_1 is not None else len(prev_pil) + meta["chunk1_num_context_latent"] = num_ctx1 + + # Align with training: when context has only 1 latent frame, context actions are omitted. + # train.py sets omit_context_actions=True when context_memory_frames == 1. + use_omit_ch1 = omit_context_actions or (num_ctx1 <= 1) + ca1 = None + if not use_omit_ch1 and num_ctx1 > 0: + ca1 = _tail_context_actions( + src_actions0, + num_ctx1, + device=device, + dtype=torch.float32, + nearest_first=(context_source == "prev_chunk_tail" and context_position == "suffix"), + ) + meta["chunk1_context_actions_count"] = int(ca1.shape[0]) if ca1 is not None else 0 + + frames_ch1 = run_one_chunk( + pipe, + prompt, + negative_prompt, + act1, + context_latents=ctx_lat_1, + num_context_frames=num_ctx1, + context_actions_t=ca1, + chunk_frames=chunk_frames, + h=h, + w=w, + seed=seed + 1, + sigma_shift=sigma_shift, + num_inference_steps=num_inference_steps, + cfg_scale=cfg_scale, + inference_noise_level=inference_noise_level, + omit_context_actions=use_omit_ch1, + context_position=context_position, + log_prefix=log_prefix + " ch1", + ) + + meta["note"] = "No cross-chunk SSM/RNN state; only frame-conditioned second chunk (same as replay eval)." + return frames_ch0, frames_ch1, meta + +def load_model(checkpoint_path, model_paths, lora_path=None, lora_alpha=1.0, device="cuda"): + """Load model from checkpoint""" + print(f"Loading model from checkpoint: {checkpoint_path}") + + # Load base pipeline + pipe = WanVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device=device, + model_configs=[ + ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="diffusion_pytorch_model*.safetensors", offload_device="cpu"), + ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth", offload_device="cpu"), + ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="Wan2.1_VAE.pth", offload_device="cpu"), + ], + ) + + # Load LoRA if specified + if lora_path and os.path.exists(lora_path): + print(f"Loading LoRA from: {lora_path}") + pipe.load_lora(pipe.dit, lora_path, alpha=lora_alpha) + + # Load checkpoint if specified + if checkpoint_path and os.path.exists(checkpoint_path): + print(f"Loading checkpoint from: {checkpoint_path}") + checkpoint = safe_load_file(checkpoint_path) + pipe.dit.load_state_dict(checkpoint, strict=False) + + pipe.enable_vram_management() + pipe.eval() + + return pipe + + +def sample_prompts_from_dataset(dataset, num_prompts=5): + """Randomly sample prompts from dataset""" + prompts = [] + dataset_size = len(dataset) + + if dataset_size == 0: + print("Warning: Dataset is empty, using default prompts") + return ["A cyberpunk city game scene, a character walking through neon-lit streets"] * num_prompts + + # Sample random indices + indices = random.sample(range(dataset_size), min(num_prompts, dataset_size)) + + print(f"Sampling {len(indices)} prompts from dataset (size: {dataset_size})...") + for idx in indices: + try: + sample = dataset[idx] + if isinstance(sample, dict): + prompt = sample.get("description") or sample.get("prompt") or sample.get("text", "") + if prompt: + prompts.append(prompt) + else: + print(f"Warning: Sample {idx} has no prompt field, skipping") + else: + print(f"Warning: Sample {idx} is not a dict, skipping") + except Exception as e: + print(f"Warning: Failed to load sample {idx}: {e}, skipping") + + # Fill with default if not enough prompts + while len(prompts) < num_prompts: + prompts.append("A cyberpunk city game scene, a character walking through neon-lit streets") + + return prompts[:num_prompts] + + +def encode_frames_to_latents(pipe, frames): + """Encode frames to latents using VAE""" + pipe.load_models_to_device(["vae"]) + vae = pipe.vae + + latents_list = [] + for frame in frames: + vid = pipe.preprocess_video([frame]).squeeze(0) + with torch.no_grad(): + lat = vae.encode([vid], device=pipe.device)[0].unsqueeze(0) + latents_list.append(lat) + + if latents_list: + return torch.cat(latents_list, dim=2) + return None + + +def generate_long_video( + pipe, + prompt, + negative_prompt="oversaturated colors, overexposed, static, blurry details", + output_dir="./long_video_output", + video_name="long_video", + context_memory_frames=4, + frames_per_segment=81, + target_frames=450, # 30 seconds at 15fps + height=352, + width=640, + num_inference_steps=20, + cfg_scale=5.0, + timestep_shift=1.0, + seed=42, + fps=15, +): + """ + Generate long video using iterative context-based generation + + Args: + pipe: WanVideoPipeline instance + prompt: Text prompt for generation + negative_prompt: Negative prompt + output_dir: Output directory for videos + video_name: Base name for output video + context_memory_frames: Number of context frames to use (K) + frames_per_segment: Frames to generate per segment (default: 81) + target_frames: Target total frames (default: 450 for 30s at 15fps) + height: Video height + width: Video width + num_inference_steps: Number of inference steps + cfg_scale: CFG scale + timestep_shift: Timestep shift + seed: Random seed + fps: FPS for output video + """ + os.makedirs(output_dir, exist_ok=True) + + # Set environment variable for concatenation inference + os.environ["USE_CONCATENATION_INFERENCE"] = "true" + + all_frames = [] + current_context_latents = None + current_context_frames = [] + + # Calculate number of segments needed + num_segments = (target_frames + frames_per_segment - 1) // frames_per_segment + + print(f"Generating long video: {target_frames} frames in {num_segments} segments") + print(f" - Frames per segment: {frames_per_segment}") + print(f" - Context frames: {context_memory_frames}") + print(f" - Prompt: {prompt[:100]}...") + + torch.manual_seed(seed) + + for segment_idx in range(num_segments): + # Calculate frames to generate for this segment + remaining_frames = target_frames - len(all_frames) + frames_to_generate = min(frames_per_segment, remaining_frames) + + if frames_to_generate <= 0: + break + + print(f"\n[{segment_idx + 1}/{num_segments}] Generating {frames_to_generate} frames...") + + # Prepare sampling kwargs + # First segment: no context (generate from scratch) + # Subsequent segments: use context from previous segment + has_context = current_context_latents is not None and segment_idx > 0 + + sampling_kwargs = { + "prompt": prompt, + "negative_prompt": negative_prompt, + "height": height, + "width": width, + "num_frames": frames_to_generate, + "num_inference_steps": num_inference_steps, + "seed": seed + segment_idx, # Different seed for each segment + "cfg_scale": cfg_scale, + "sigma_shift": timestep_shift, + "denoising_strength": 1.0, + } + + # Add context memory only if we have context + if has_context: + sampling_kwargs["enable_context_memory"] = True + sampling_kwargs["context_latents"] = current_context_latents + sampling_kwargs["num_context_frames"] = len(current_context_frames) + + try: + # Generate frames + if has_context: + print(f" Using {len(current_context_frames)} context frames from previous segment...") + + generated_frames = pipe(**sampling_kwargs) + + if isinstance(generated_frames, list): + segment_frames = generated_frames + else: + segment_frames = [generated_frames] if hasattr(generated_frames, '__iter__') else [generated_frames] + + # Add to all frames + all_frames.extend(segment_frames) + + # Update context: use last K frames from generated segment + # These will be used as context for the next segment + if len(segment_frames) >= context_memory_frames: + context_frames = segment_frames[-context_memory_frames:] + current_context_frames = context_frames + + # Encode context frames to latents + print(f" Encoding last {context_memory_frames} frames as context for next segment...") + current_context_latents = encode_frames_to_latents(pipe, context_frames) + else: + # If not enough frames, use all frames as context + current_context_frames = segment_frames + current_context_latents = encode_frames_to_latents(pipe, segment_frames) + + print(f" Generated {len(segment_frames)} frames (total: {len(all_frames)}/{target_frames})") + + except Exception as e: + print(f" Error generating segment {segment_idx + 1}: {e}") + traceback.print_exc() + break + + # Save final video + if len(all_frames) > 0: + output_path = os.path.join(output_dir, f"{video_name}.mp4") + print(f"\nSaving video to: {output_path}") + print(f" Total frames: {len(all_frames)}") + print(f" Duration: {len(all_frames) / fps:.2f} seconds") + + save_video(all_frames, output_path, fps=fps, quality=5) + print(f"Video saved: {output_path}") + + # Save prompt + prompt_path = os.path.join(output_dir, f"{video_name}_prompt.txt") + with open(prompt_path, 'w', encoding='utf-8') as f: + f.write(prompt) + + return output_path + else: + print("Error: No frames generated") + return None + + +def main(): + parser = argparse.ArgumentParser(description="Generate long videos using iterative context-based generation") + + # Model paths + parser.add_argument("--checkpoint_path", type=str, default=None, help="Path to model checkpoint") + parser.add_argument("--lora_path", type=str, default=None, help="Path to LoRA weights") + parser.add_argument("--lora_alpha", type=float, default=1.0, help="LoRA alpha") + parser.add_argument("--model_paths", type=str, default=None, help="JSON string of model paths (not used if checkpoint_path is set)") + + # Dataset + parser.add_argument("--dataset_base_path", type=str, required=True, help="Base path to dataset") + parser.add_argument("--dataset_metadata_path", type=str, required=True, help="Path to dataset metadata CSV") + parser.add_argument("--num_prompts", type=int, default=5, help="Number of prompts to sample from dataset") + + # Generation parameters + parser.add_argument("--output_dir", type=str, default="./long_video_output", help="Output directory") + parser.add_argument("--context_memory_frames", type=int, default=4, help="Number of context frames (K)") + parser.add_argument("--frames_per_segment", type=int, default=81, help="Frames per segment (default: 81)") + parser.add_argument("--target_frames", type=int, default=450, help="Target total frames (30s at 15fps)") + parser.add_argument("--height", type=int, default=352, help="Video height") + parser.add_argument("--width", type=int, default=640, help="Video width") + parser.add_argument("--num_inference_steps", type=int, default=20, help="Number of inference steps") + parser.add_argument("--cfg_scale", type=float, default=5.0, help="CFG scale") + parser.add_argument("--timestep_shift", type=float, default=1.0, help="Timestep shift") + parser.add_argument("--seed", type=int, default=42, help="Random seed") + parser.add_argument("--fps", type=int, default=15, help="FPS for output video") + parser.add_argument("--device", type=str, default="cuda", help="Device (cuda/cpu)") + + args = parser.parse_args() + + # Load dataset for prompt sampling + print("Loading dataset...") + dataset_args = wan_parser.parse_args([]) # Create minimal args + dataset_args.dataset_base_path = args.dataset_base_path + dataset_args.dataset_metadata_path = args.dataset_metadata_path + dataset_args.height = args.height + dataset_args.width = args.width + + dataset = VideoDataset(args=dataset_args) + print(f"Dataset loaded: {len(dataset)} samples") + + # Sample prompts + prompts = sample_prompts_from_dataset(dataset, args.num_prompts) + print(f"Sampled {len(prompts)} prompts") + + # Load model + model_paths = None + if args.model_paths: + model_paths = json.loads(args.model_paths) + + pipe = load_model( + checkpoint_path=args.checkpoint_path, + model_paths=model_paths, + lora_path=args.lora_path, + lora_alpha=args.lora_alpha, + device=args.device, + ) + + # Generate videos for each prompt + output_paths = [] + for idx, prompt in enumerate(prompts): + print(f"\n{'='*80}") + print(f"Generating video {idx + 1}/{len(prompts)}") + print(f"{'='*80}") + + video_name = f"long_video_{idx + 1:03d}" + + output_path = generate_long_video( + pipe=pipe, + prompt=prompt, + output_dir=args.output_dir, + video_name=video_name, + context_memory_frames=args.context_memory_frames, + frames_per_segment=args.frames_per_segment, + target_frames=args.target_frames, + height=args.height, + width=args.width, + num_inference_steps=args.num_inference_steps, + cfg_scale=args.cfg_scale, + timestep_shift=args.timestep_shift, + seed=args.seed + idx, # Different seed for each video + fps=args.fps, + ) + + if output_path: + output_paths.append(output_path) + + print(f"\n{'='*80}") + print(f"Generation completed: {len(output_paths)} videos generated") + print(f"Output directory: {args.output_dir}") + print(f"{'='*80}") diff --git a/code/src/model_training/train.py b/code/src/model_training/train.py new file mode 100644 index 0000000000000000000000000000000000000000..933fd8817c38a7283d35a8845766a51dd25e039d --- /dev/null +++ b/code/src/model_training/train.py @@ -0,0 +1,661 @@ +import os, sys, re +import torch +import torch.nn as nn +import logging + +logger = logging.getLogger(__name__) + +_rank_env = os.environ.get("RANK") or os.environ.get("LOCAL_RANK") or os.environ.get("ACCELERATE_PROCESS_INDEX") or "0" +_rank = int(str(_rank_env)) +_level = logging.INFO if _rank == 0 else logging.WARNING +logging.basicConfig( + level=_level, + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + force=True, +) +logger.setLevel(_level) + +current_file_abs = os.path.abspath(__file__) +project_root = os.path.dirname(os.path.dirname(os.path.dirname(current_file_abs))) + +if project_root not in sys.path: + sys.path.insert(0, project_root) + +modules_to_clear = [ + 'diffsynth.models.memory.framepack_length', + 'diffsynth.models.memory.framepack_weight', + 'diffsynth.models.memory.spatial_grid_memory', + 'diffsynth.models.memory.videossm_hybrid', + 'diffsynth.models.memory.block_wise_ssm', + 'diffsynth.models.memory', + 'diffsynth.pipelines.wan_video_new', + 'diffsynth.trainers.utils', + 'diffsynth.models.wan_video_dit', + 'diffsynth.lora.flux_lora', + 'diffsynth.lora', + 'diffsynth.configs.model_config', + 'diffsynth.configs', + 'diffsynth.pipelines', + 'diffsynth.trainers', + 'diffsynth.models', + 'diffsynth', +] + +for mod in modules_to_clear: + if mod in sys.modules: + del sys.modules[mod] + +import importlib +importlib.invalidate_caches() + +from diffsynth.pipelines.wan_video_new import WanVideoPipeline, ModelConfig + +try: + import transformers + if not hasattr(transformers, "HybridCache") and hasattr(transformers, "DynamicCache"): + transformers.HybridCache = transformers.DynamicCache +except Exception: + pass + +from diffsynth.trainers.utils import DiffusionTrainingModule, ModelLogger as BaseModelLogger, VideoDataset, CamVideoDataset, wan_parser +from diffsynth.models.wan_video_dit import SelfAttention, CrossAttention, GateModule, modulate +from diffsynth.models.memory.videossm_hybrid import HybridStateSpaceMemory +from diffsynth.models.memory.block_wise_ssm import BlockWiseStateSpaceMemory +from diffsynth.models.memory.u_vit_cgla_blocks import CGLATransformerBlock, remap_wan_to_cgla + +try: + import diffsynth.trainers.utils as utils_module + utils_file = utils_module.__file__ if hasattr(utils_module, '__file__') else 'unknown' + is_local = 'site-packages' not in utils_file + if is_local: + logger.info(f"[VERIFIED] Using LOCAL diffsynth code from: {utils_file}") + else: + logger.warning(f"Using INSTALLED diffsynth package from: {utils_file}") +except Exception as e: + logger.error(f"Failed to verify code location: {e}") + +import random +import numpy as np +os.environ["TOKENIZERS_PARALLELISM"] = "false" +from safetensors.torch import load_file as safe_load_file +from src.model_training.fov_retrieval import setup_fov_retriever_for_training +from src.model_training.training_modules import DiTBlock_w_Action, WanTrainingModule + + +def set_seed(seed=42): + """Set random seeds for reproducible training.""" + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + os.environ['PYTHONHASHSEED'] = str(seed) + os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8' + logger.info(f"Random seed set to {seed}") + + +def _log_dit_freeze_summary(dit: torch.nn.Module) -> None: + by_module: dict[str, tuple[int, bool]] = {} + for name, p in dit.named_parameters(): + numel = p.numel() + trainable = p.requires_grad + parts = name.split(".") + prefix = ".".join(parts[:-1]) if len(parts) > 1 else name + if prefix not in by_module: + by_module[prefix] = (0, False) + prev_numel, prev_trainable = by_module[prefix] + by_module[prefix] = (prev_numel + numel, prev_trainable or trainable) + trainable_list = [(k, v[0]) for k, v in by_module.items() if v[1]] + frozen_list = [(k, v[0]) for k, v in by_module.items() if not v[1]] + trainable_list.sort(key=lambda x: x[0]) + frozen_list.sort(key=lambda x: x[0]) + total_trainable = sum(n for _, n in trainable_list) + total_frozen = sum(n for _, n in frozen_list) + examples = ", ".join(name for name, _ in trainable_list[:8]) + logger.info( + f"[DiT freeze] trainable={total_trainable:,} ({len(trainable_list)} groups), " + f"frozen={total_frozen:,} ({len(frozen_list)} groups), examples=[{examples}]" + ) + + +set_seed(42) + + +from src.model_training.training_modules.model_logger import ModelLogger +from src.model_training.training_modules.training_loop import launch_training_task + + +if __name__ == "__main__": + parser = wan_parser() + def _add_arg_if_missing(*args, **kwargs): + if args and args[0] in parser._option_string_actions: + return + parser.add_argument(*args, **kwargs) + + for name, kwargs in [ + ("--tokenizer_path", dict(type=str, default=None, help="Local tokenizer path.")), + ("--wandb_run_name", dict(type=str, default=None)), + ("--ckpt_interval", dict(type=int, default=None)), + ("--trainable_dit_modules", dict(type=str, default=None, help="Comma-separated DiT modules to unfreeze.")), + ("--num_workers", dict(type=int, default=0, help="DataLoader workers.")), + ("--max_train_steps", dict(type=int, default=0, help="Stop after N optimizer steps.")), + ("--progress_total_steps", dict(type=int, default=0, help="tqdm total steps override.")), + ("--log_interval", dict(type=int, default=20, help="Log loss/grad_norm/lr every N optimizer steps.")), + ("--max_grad_norm", dict(type=float, default=1.0, help="Grad-norm clip (default 1.0; 0 = no clip, norm still logged).")), + ("--resume_from_checkpoint", dict(type=str, default=None)), + ("--context_memory_frames", dict(type=int, default=8)), + ("--training_mode", dict(type=str, default="predict", choices=["predict", "context", "condition"])), + ("--context_drop_prob", dict(type=float, default=0.0)), + ("--retrieval_method", dict(type=str, default="fov", choices=["fov", "latent_sim"])), + ("--latent_retrieval_dir", dict(type=str, default=None)), + ("--fov_top_k", dict(type=int, default=4)), + ("--context_attention_weight", dict(type=float, default=1.0)), + ("--context_temporal_decay", dict(type=float, default=1.0)), + ("--spike_threshold", dict(type=float, default=5.0)), + ("--spatial_memory_tokens", dict(type=int, default=64)), + ("--spatial_memory_grid", dict(type=int, default=8)), + ("--spatial_memory_inject_mode", dict(type=str, default="concat_text", choices=["concat_text", "none", "cross_attn_readout"])), + ("--geometry_memory_column", dict(type=str, default="geometry_memory")), + ("--geometry_memory_root", dict(type=str, default=None)), + ("--geometry_spatial_memory_tokens", dict(type=int, default=64)), + ("--geometry_spatial_memory_grid", dict(type=int, default=8)), + ("--geometry_spatial_memory_temporal_bins", dict(type=int, default=4)), + ("--geometry_spatial_memory_inject_mode", dict(type=str, default="concat_text", choices=["concat_text", "none", "cross_attn_readout"])), + ("--framepack_ratio", dict(type=int, default=2)), + ("--framepack_length_strategy", dict(type=str, default="distance_merge", choices=["distance_merge", "mean", "uniform", "recent_weighted", "weighted_recent", "packed_multiscale"])), + ("--framepack_recent_keep_ratio", dict(type=float, default=0.5)), + ("--framepack_multiscale_w2", dict(type=float, default=0.25)), + ("--framepack_multiscale_w4", dict(type=float, default=0.15)), + ("--context_source", dict(type=str, default="fov", choices=["fov", "replay", "prev_chunk_tail"])), + ("--ssm_num_blocks_hint", dict(type=int, default=21)), + ("--ssm_every_n_blocks", dict(type=int, default=4)), + ("--videossm_kernel_size", dict(type=int, default=3)), + ("--videossm_expand", dict(type=int, default=2)), + ("--videossm_every_n_blocks", dict(type=int, default=4)), + # Camera-Guided Linear Attention (CGLA) memory row. + ("--cgla_every_n_blocks", dict(type=int, default=4)), + ("--cgla_num_heads", dict(type=int, default=0)), # 0 -> auto (dim // head_dim) + ("--cgla_head_dim", dict(type=int, default=128)), + ("--cgla_num_sparse_partition", dict(type=int, default=4)), + ("--cgla_num_writer", dict(type=int, default=1)), + ("--cgla_num_reader", dict(type=int, default=1)), + ("--cgla_gate_logit_normalizer", dict(type=int, default=16)), + ("--cgla_gate_low_rank_dim", dict(type=int, default=16)), + ("--cgla_pose_dim", dict(type=int, default=12)), + ("--cgla_pose_bottleneck", dict(type=int, default=64)), + ("--cgla_aux_loss_weight", dict(type=float, default=0.0)), + ("--cgla_mechanism", dict(type=str, default="cgla", + choices=["cgla", "prope", "ucpe"])), + ("--cgla_emb_dim", dict(type=int, default=1024)), + ("--cgla_num_patches", dict(type=int, default=880)), + ("--cgla_temporal_length", dict(type=int, default=21)), + ("--cgla_ffn_dim", dict(type=int, default=0)), # 0 -> use Wan DiT ffn_dim + ("--sampling_interval_steps", dict(type=int, default=0)), + ("--sampling_negative_prompt", dict(type=str, default="oversaturated colors, overexposed, static, blurry details")), + ("--sampling_height", dict(type=int, default=352)), + ("--sampling_width", dict(type=int, default=640)), + ("--sampling_num_frames", dict(type=int, default=81)), + ("--sampling_num_inference_steps", dict(type=int, default=50)), + ("--sampling_action_path", dict(type=str, default=None)), + ("--sampling_two_chunk_action_path", dict(type=str, default=None)), + ("--sampling_eval_dataset_base", dict(type=str, default=None)), + ("--sampling_eval_metadata_path", dict(type=str, default=None)), + ("--samples_per_epoch", dict(type=int, default=0)), + ("--camera_encoder_scale", dict(type=float, default=1.0)), + ("--camera_inject_mode", dict(type=str, default="post", choices=["post", "pre_norm", "pre_qkv", "pre_qkv_post", "pre_modulate", "pre_qkv_gated"])), + ]: + _add_arg_if_missing(name, **kwargs) + + for name in [ + "--save_full_model", "--disable_gradient_checkpointing", "--add_action_attn", "--action_use_temporal_attention", + "--action_inject_after_spatial_attn", "--use_camera_encoder", "--camera_encoder_shallow", + "--camera_encoder_separate_t_r", "--camera_encoder_explicit_yaw", "--yaw_flip_aug", + "--camera_encoder_sincos_yaw", "--camera_encoder_r_mlp_no_layernorm", + "--add_camera_outside_gate", "--no_camera_encoder_zero_init", + "--camera_encoder_full_zero_init", "--enable_context_memory", "--context_per_frame_vae", + "--cfg_target_only", "--enable_fov_retrieval", "--use_rt_relative", + "--strict_overlap_context", "--use_anchor_frame", "--use_spatial_memory", + "--use_geometry_spatial_memory", + "--use_spatial_memory_legacy", "--use_framepack_memory", "--use_framepack_length_compress", + "--use_block_wise_ssm", "--use_videossm_hybrid", "--sampling_two_chunk_memory", + "--use_cgla_memory", "--cgla_use_pose_rope", "--cgla_use_pose_gate_mod", + ]: + _add_arg_if_missing(name, action="store_true") + + for name, kwargs in [ + ("--per_device_train_batch_size", dict(type=int, default=None)), + ("--timestep_shift", dict(type=float, default=1.0)), + ("--action_base_path", dict(type=str, default=None)), + ("--ckpt_path", dict(type=str, default=None)), + ("--cam_position_scale", dict(type=float, default=0.01)), + ("--resume_from", dict(type=str, default=None)), + ("--verify_ckpt_step", dict(type=int, default=0)), + ("--verify_high_noise_first_steps", dict(type=int, default=0)), + ("--moc_temperature", dict(type=float, default=1.0)), + ("--moc_top_k", dict(type=int, default=0)), + ("--prev_chunk_frames", dict(type=int, default=81)), + ("--implicit_type", dict(type=str, default="summary")), + ("--context_compressor_ratio", dict(type=int, default=2)), + ("--episodic_buffer_size", dict(type=int, default=0)), + ("--episodic_replay_interval", dict(type=int, default=0)), + ("--episodic_replay_weight", dict(type=float, default=0.0)), + ]: + _add_arg_if_missing(name, **kwargs) + for name in [ + "--enable_video_sampling", "--sampling_atomic_left_right", "--sampling_four_prompts", + "--sampling_two_prompts", "--train_action_module", "--train_cam_pose", + "--action_module_only", "--use_moc", "--unified_implicit", "--use_implicit_memory", + "--use_memory_v2v_compressor", "--use_slow_fast_memory", "--use_entity_memory", + "--use_episodic_memory", + ]: + _add_arg_if_missing(name, action="store_true") + args = parser.parse_args() + def _arg(name, default=None): + return getattr(args, name, default) + + def _normalize_and_validate_args(): + # Backward-compat mappings + if _arg("per_device_train_batch_size", None) is None: + args.per_device_train_batch_size = int(_arg("batch_size", 1)) + if _arg("sampling_atomic_left_right", False) and not _arg("sampling_two_chunk_memory", False): + # Legacy monitor intent maps to current two-chunk monitor. + args.sampling_two_chunk_memory = True + if _arg("enable_video_sampling", False) and int(_arg("sampling_interval_steps", 0)) <= 0: + args.sampling_interval_steps = 1000 + + # Keep paper-style block-wise SSM and legacy VideoSSM hybrid explicitly separated. + if _arg("use_block_wise_ssm", False) and _arg("use_videossm_hybrid", False): + raise ValueError( + "--use_block_wise_ssm and --use_videossm_hybrid are mutually exclusive; " + "use block-wise SSM for paper-aligned runs or VideoSSM hybrid for legacy baselines." + ) + + # CGLA is another per-block temporal-memory module on the same hook as + # block-wise SSM / VideoSSM hybrid; only one may be active at a time. + _per_block_memory_active = ( + bool(_arg("use_block_wise_ssm", False)) + or bool(_arg("use_videossm_hybrid", False)) + ) + if _arg("use_cgla_memory", False) and _per_block_memory_active: + raise ValueError( + "--use_cgla_memory is mutually exclusive with --use_block_wise_ssm / " + "--use_videossm_hybrid (two per-block temporal-memory modules on the same hook)." + ) + if _arg("use_cgla_memory", False) and not _arg("train_cam_pose", False): + logger.warning( + "--use_cgla_memory expects per-frame camera pose; enabling without " + "--train_cam_pose means no RT reaches the module (it will run as vanilla GLA)." + ) + + # Explicit retrieval strategy visibility: default fov, latent_sim degrades to fov when cache dir is absent. + if _arg("retrieval_method", "fov") == "latent_sim": + if not _arg("latent_retrieval_dir", None): + logger.warning("retrieval_method=latent_sim but latent_retrieval_dir is empty; runtime will fallback to fov retrieval.") + else: + logger.info(f"retrieval_method=latent_sim latent_retrieval_dir={args.latent_retrieval_dir}") + else: + logger.info("retrieval_method=fov") + + # 2-chunk sampling defaults: keep left/right_45 semantics compatible with existing shell wrappers. + if _arg("sampling_two_chunk_action_path", None) in (None, ""): + args.sampling_two_chunk_action_path = _arg("sampling_action_path", None) + + _normalize_and_validate_args() + + resume_step_count = 0 + if args.resume_from_checkpoint is not None: + if (_arg('trainable_dit_modules', None) or "").strip() or _arg('resume_weights_only', False): + logger.info("resume_from_checkpoint used for weights only (trainable_dit_modules set or resume_weights_only), step count starts from 0, no skip data") + resume_step_count = 0 + else: + checkpoint_filename = os.path.basename(args.resume_from_checkpoint) + step_match = re.search(r'Step-(\d+)', checkpoint_filename) + epoch_match = re.search(r'epoch-(\d+)', checkpoint_filename) + if step_match: + resume_step_count = int(step_match.group(1)) + logger.info(f"Resuming from step {resume_step_count} (extracted from checkpoint filename)") + elif epoch_match: + logger.info(f"Resuming from epoch checkpoint (epoch-{epoch_match.group(1)}), step count will start from 0") + resume_step_count = 0 + else: + logger.warning("Could not extract step count from checkpoint filename, starting from step 0") + + set_seed(42) + + args.enable_icl = False + args.icl_num_examples = 2 + args.icl_context_frames = 8 + + if _arg('train_cam_pose', False): + dataset = CamVideoDataset(args=args) + else: + dataset = VideoDataset(args=args, action_base_path=args.action_base_path) + + def _log_dataset_validation(ds): + ds_size = len(ds) + ds_repeat = _arg('dataset_repeat', 1) + logger.info( + f"[Dataset] size={ds_size}, repeat={ds_repeat}, " + f"epochs={args.num_epochs}, total_samples={ds_size * ds_repeat * args.num_epochs}" + ) + + _log_dataset_validation(dataset) + + model = WanTrainingModule( + model_paths=args.model_paths, + model_id_with_origin_paths=args.model_id_with_origin_paths, + tokenizer_path=_arg('tokenizer_path', None), + trainable_models=_arg('trainable_models', None), + lora_base_model=args.lora_base_model, + lora_target_modules=args.lora_target_modules, + lora_rank=args.lora_rank, + use_gradient_checkpointing=not _arg('disable_gradient_checkpointing', False), + use_gradient_checkpointing_offload=args.use_gradient_checkpointing_offload, + extra_inputs=args.extra_inputs, + resume_from_checkpoint=args.resume_from_checkpoint, + dataset_base_path=_arg('dataset_base_path', None), + enable_context_memory=_arg('enable_context_memory', False), + context_drop_prob=_arg('context_drop_prob', 0.0), + context_drop_seed=42, + omit_context_actions=_arg('omit_context_actions', False) or (_arg('context_memory_frames', 8) == 1), # ctx=1: no context action injection + context_noise_prob=_arg('context_noise_prob', 0.0), + context_noise_std=_arg('context_noise_std', 0.02), + context_fixed_noise_std=_arg('context_fixed_noise_std', None), + context_memory_frames=_arg('context_memory_frames', 8), + context_per_frame_vae=_arg('context_per_frame_vae', False), + training_mode=_arg('training_mode', 'predict'), + teacher_forcing_prob=_arg('teacher_forcing_prob', 0.0), + yaw_flip_aug=_arg('yaw_flip_aug', False), + context_source=_arg('context_source', 'fov'), + use_framepack_memory=_arg('use_framepack_memory', False), + context_temporal_decay=_arg('context_temporal_decay', 1.0), + context_attention_weight=_arg('context_attention_weight', 1.0), + use_framepack_length_compress=_arg('use_framepack_length_compress', False), + framepack_ratio=_arg('framepack_ratio', 2), + framepack_length_strategy=_arg('framepack_length_strategy', 'distance_merge'), + framepack_recent_keep_ratio=_arg('framepack_recent_keep_ratio', 0.5), + framepack_multiscale_w2=_arg('framepack_multiscale_w2', 0.25), + framepack_multiscale_w4=_arg('framepack_multiscale_w4', 0.15), + use_spatial_memory=_arg('use_spatial_memory', False), + use_spatial_memory_legacy=_arg('use_spatial_memory_legacy', False), + spatial_memory_tokens=_arg('spatial_memory_tokens', 64), + spatial_memory_grid=_arg('spatial_memory_grid', 8), + spatial_memory_inject_mode=_arg('spatial_memory_inject_mode', 'concat_text'), + use_geometry_spatial_memory=_arg('use_geometry_spatial_memory', False), + geometry_spatial_memory_tokens=_arg('geometry_spatial_memory_tokens', 64), + geometry_spatial_memory_grid=_arg('geometry_spatial_memory_grid', 8), + geometry_spatial_memory_temporal_bins=_arg('geometry_spatial_memory_temporal_bins', 4), + geometry_spatial_memory_inject_mode=_arg( + 'geometry_spatial_memory_inject_mode', + 'concat_text', + ), + use_moc=_arg('use_moc', False), + moc_temperature=_arg('moc_temperature', 1.0), + moc_top_k=_arg('moc_top_k', 0), + timestep_shift=float(_arg('timestep_shift', 1.0)), + ) + if _arg('use_moc', False): + logger.info( + f"[MoC] enabled with temperature={float(_arg('moc_temperature', 1.0))}, " + f"top_k={int(_arg('moc_top_k', 0) or 0)}" + ) + if _arg('use_geometry_spatial_memory', False): + logger.info( + "[Geometry Spatial Memory] enabled; expects TSDF/point-cloud rendered condition " + f"from metadata column '{_arg('geometry_memory_column', 'geometry_memory')}'" + ) + + # ── VWM-style: Replace DiT blocks with DiTBlock_w_Action ── + _use_cam_pose = bool(_arg('train_cam_pose', False)) + use_cgla_memory = bool(_arg('use_cgla_memory', False)) + if _arg('train_action_module', False) or _use_cam_pose or use_cgla_memory: + dit = model.pipe.dit + old_blocks = dit.blocks + has_image_input = dit.has_image_input + dim = dit.dim + num_heads = dit.num_heads + ffn_dim = dit.ffn_dim + eps = 1e-6 + + block_dtype = next(old_blocks[0].parameters()).dtype + + use_block_wise_ssm = bool(_arg('use_block_wise_ssm', False)) + use_videossm_hybrid = bool(_arg('use_videossm_hybrid', False)) + ssm_every_n = max(int(_arg('ssm_every_n_blocks', 4)), 1) + videossm_every_n = max(int(_arg('videossm_every_n_blocks', 4)), 1) + # CGLA: which DiT blocks become CGLATransformerBlock (the rest stay + # DiTBlock_w_Action = Wan softmax + cam-pose). cgla_every_n_blocks=1 => + # all blocks (matches tests/test_cgla_wan.py); =4 => every 4th block, + # the same attach cadence as the SSM/VideoSSM rows (controlled ablation). + cgla_every_n = max(int(_arg('cgla_every_n_blocks', 4) or 4), 1) + _cgla_head_dim = int(_arg('cgla_head_dim', 128) or 128) + _cgla_head_dim = (_cgla_head_dim if _cgla_head_dim > 0 else (dim // num_heads)) + + new_blocks = nn.ModuleList() + for block_id, old_block in enumerate(old_blocks): + attach_block_ssm = use_block_wise_ssm and (block_id % ssm_every_n == 0) + attach_videossm = use_videossm_hybrid and (block_id % videossm_every_n == 0) + attach_cgla = use_cgla_memory and (block_id % cgla_every_n == 0) + if attach_cgla: + # CGLA row: the DiT block IS CGLATransformerBlock (Wan DiTBlock + # with self-attn swapped for the SSE-GLA). Wan's self_attn.{q,k,v,o} + # initialise the SSE-GLA {q,k,v,o}_proj (same shapes, via remap); + # cross_attn/norm/ffn/modulation load directly. + new_block = CGLATransformerBlock( + has_image_input=has_image_input, + dim=dim, num_heads=num_heads, ffn_dim=ffn_dim, eps=eps, + head_dim=_cgla_head_dim, + num_sparse_partition=int(_arg('cgla_num_sparse_partition', 4) or 4), + num_writer=int(_arg('cgla_num_writer', 1) or 1), + num_reader=int(_arg('cgla_num_reader', 1) or 1), + pose_dim=int(_arg('cgla_pose_dim', 12) or 12), + pose_bottleneck=int(_arg('cgla_pose_bottleneck', 64) or 64), + gate_logit_normalizer=int(_arg('cgla_gate_logit_normalizer', 16) or 16), + gate_low_rank_dim=int(_arg('cgla_gate_low_rank_dim', 16) or 16), + use_pose_gate_mod=bool(_arg('cgla_use_pose_gate_mod', False)), + layer_idx=block_id, + mechanism=str(_arg('cgla_mechanism', 'cgla') or 'cgla'), + bidirectional=True, + add_action_attn=_arg('add_action_attn', False), + action_use_temporal_attention=_arg('action_use_temporal_attention', False), + use_cam_pose=_use_cam_pose, + ) + new_block = new_block.to(dtype=block_dtype, device=next(old_block.parameters()).device) + # Copy the Wan submodules that match by name (cross_attn/norm/ffn). + for attr in ("cross_attn", "norm1", "norm2", "norm3", "ffn"): + if hasattr(old_block, attr) and hasattr(new_block, attr): + getattr(new_block, attr).load_state_dict(getattr(old_block, attr).state_dict()) + if hasattr(old_block, "modulation") and hasattr(new_block, "modulation"): + with torch.no_grad(): + new_block.modulation.copy_(old_block.modulation.to(dtype=block_dtype)) + # # self_attn: SSE-GLA structure differs from Wan SelfAttention, + # # so remap Wan's {q,k,v,o} -> {q,k,v,o}_proj and load (norm_q/ + # # norm_k have no SSE-GLA counterpart -> dropped). + # _sa_sd = {} + # for _k, _v in old_block.self_attn.state_dict().items(): + # if _k.startswith("q."): + # _sa_sd["q_proj." + _k[2:]] = _v + # elif _k.startswith("k."): + # _sa_sd["k_proj." + _k[2:]] = _v + # elif _k.startswith("v."): + # _sa_sd["v_proj." + _k[2:]] = _v + # elif _k.startswith("o."): + # _sa_sd["o_proj." + _k[2:]] = _v + # new_block.self_attn.load_state_dict(_sa_sd, strict=False) + else: + new_block = DiTBlock_w_Action( + has_image_input=has_image_input, + dim=dim, num_heads=num_heads, ffn_dim=ffn_dim, eps=eps, + add_action_attn=_arg('add_action_attn', False), + action_use_temporal_attention=_arg('action_use_temporal_attention', False), + use_cam_pose=_use_cam_pose, + use_block_wise_ssm=attach_block_ssm, + use_videossm_hybrid=attach_videossm, + videossm_kernel_size=int(_arg('videossm_kernel_size', 3) or 3), + videossm_expand=int(_arg('videossm_expand', 2) or 2), + ) + new_block = new_block.to(dtype=block_dtype, device=next(old_block.parameters()).device) + for attr in ("self_attn", "cross_attn", "norm1", "norm2", "norm3", "ffn"): + if hasattr(old_block, attr) and hasattr(new_block, attr): + getattr(new_block, attr).load_state_dict(getattr(old_block, attr).state_dict()) + if hasattr(old_block, "modulation") and hasattr(new_block, "modulation"): + with torch.no_grad(): + new_block.modulation.copy_(old_block.modulation.to(dtype=block_dtype)) + new_blocks.append(new_block) + + dit.blocks = new_blocks + _n_cgla = sum(1 for b in new_blocks if isinstance(b, CGLATransformerBlock)) + if use_cgla_memory: + _mech = str(_arg('cgla_mechanism', 'cgla') or 'cgla') + logger.info(f"[CGLA] Replaced {_n_cgla}/{len(new_blocks)} DiT blocks with " + f"CGLATransformerBlock (every_n={cgla_every_n}, mechanism={_mech}, " + f"head_dim={_cgla_head_dim}, bidirectional); " + f"Wan self_attn q/k/v/o -> SSE-GLA q/k/v/o_proj") + + _mlp_type = "MLP_CamPose" if _use_cam_pose else "MLP_Action" + logger.info(f"[VWM-style] Replaced {len(new_blocks)} DiT blocks with DiTBlock_w_Action ({_mlp_type}, zero-init)") + if use_block_wise_ssm: + logger.info(f"[Block-wise SSM] attached to every {ssm_every_n} DiT block(s)") + if use_videossm_hybrid: + logger.info(f"[VideoSSM hybrid] attached to every {videossm_every_n} DiT block(s)") + + device = next(dit.parameters()).device + _ckpt_path = _arg('ckpt_path', None) or _arg('resume_from_checkpoint', None) + if _ckpt_path is not None: + ckpt = safe_load_file(_ckpt_path) + if use_cgla_memory: + # Remap Wan self_attn.{q,k,v,o} -> SSE-GLA {q,k,v,o}_proj so a Wan + # checkpoint loads the linear-attention projections. No-op on an + # already-CGLA checkpoint (keys are already *_proj). + ckpt = remap_wan_to_cgla(ckpt) + missing, unexpected = dit.load_state_dict(ckpt, strict=False) + dit.to(device=device) + logger.info(f"[VWM-style] Loaded ckpt: {len(ckpt)} keys, missing={len(missing)}, unexpected={len(unexpected)}") + + if use_cgla_memory: + # Block-aware freeze (supports cgla_every_n_blocks > 1, where some + # blocks stay DiTBlock_w_Action = Wan softmax). CGLA blocks: train + # the SSE-GLA self_attn + cam_encoder + cgla_gate + noise_write_gate, + # freeze the Wan backbone (cross_attn/norm/ffn/modulation). Non-CGLA + # blocks: VWM pattern (train action_mlp / self_attn_with_action / + # SSM, freeze Wan softmax self_attn + backbone). + for block in dit.blocks: + if isinstance(block, CGLATransformerBlock): + for name, param in block.named_parameters(): + if ("action_mlp" in name) or ("self_attn_with_action" in name) or ("self_attn" in name) or \ + ("cam_encoder" in name) or ("noise_write_gate" in name) or ("cgla_gate" in name): + param.requires_grad = True + else: + param.requires_grad = False + else: + for name, param in block.named_parameters(): + if ("action_mlp" in name) or ("self_attn_with_action" in name) or \ + ("block_wise_ssm" in name) or ("videossm_hybrid" in name): + param.requires_grad = True + else: + param.requires_grad = False + elif _arg('action_module_only', False): + if _arg('add_action_attn', False): + for block in dit.blocks: + for name, param in block.named_parameters(): + if ("action_mlp" in name) or ("self_attn_with_action" in name) or ("block_wise_ssm" in name) or ("videossm_hybrid" in name): + param.requires_grad = True + else: + param.requires_grad = False + else: + for block in dit.blocks: + for name, param in block.named_parameters(): + if "action_mlp" in name or "self_attn" in name or "block_wise_ssm" in name or "videossm_hybrid" in name: + param.requires_grad = True + else: + param.requires_grad = False + else: + for block in dit.blocks: + for name, param in block.named_parameters(): + if "action_mlp" in name or "self_attn_with_action" in name or "block_wise_ssm" in name or "videossm_hybrid" in name: + param.requires_grad = True + else: + param.requires_grad = False + _log_dit_freeze_summary(dit) + + _resume_from = _arg('resume_from', None) + if _resume_from: + logger.info(f"Loading full resume checkpoint: {_resume_from}") + ckpt = safe_load_file(_resume_from) + if use_cgla_memory: + # No-op on an already-CGLA ckpt (keys are *_proj); remaps a Wan ckpt. + ckpt = remap_wan_to_cgla(ckpt) + model.pipe.dit.load_state_dict(ckpt, strict=False) + logger.info(f"Checkpoint loaded, resuming from step {resume_step_count}") + + model_logger = ModelLogger( + args.output_path, + remove_prefix_in_ckpt=args.remove_prefix_in_ckpt, + wandb_run_name=args.wandb_run_name, + ckpt_interval=args.ckpt_interval, + resume_step_count=resume_step_count, + save_full_model=_arg('save_full_model', False), + context_drop_prob=float(_arg("context_drop_prob", 0.0)), + enable_video_sampling=_arg("enable_video_sampling", False), + sampling_interval_steps=int(_arg("sampling_interval_steps", 0)), + sampling_two_chunk_memory=_arg("sampling_two_chunk_memory", False), + sampling_action_path=_arg("sampling_action_path", None), + sampling_two_chunk_action_path=_arg("sampling_two_chunk_action_path", None), + sampling_negative_prompt=_arg("sampling_negative_prompt", ""), + sampling_height=int(_arg("sampling_height", 352)), + sampling_width=int(_arg("sampling_width", 640)), + sampling_num_frames=int(_arg("sampling_num_frames", 81)), + sampling_num_inference_steps=int(_arg("sampling_num_inference_steps", 50)), + context_memory_frames=int(_arg("context_memory_frames", 1)), + context_source=_arg("context_source", "replay"), + context_per_frame_vae=_arg("context_per_frame_vae", False), + # Monitor samples with the SAME noise-schedule shift training uses + # (--timestep_shift), so the in-training video reflects the trained + # schedule (15 for two-chunk rows, 5 for legacy ctx rows). + sampling_sigma_shift=float(_arg("timestep_shift", 1.0) or 1.0), + ) + + optimizer = torch.optim.AdamW(model.trainable_modules(), lr=args.learning_rate) + scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer) + + # Setup FOV retriever for context-based memory training (also for ModelLogger sampling) + enable_fov_retrieval = _arg('enable_fov_retrieval', False) + fov_retriever = None + dataset_base_path = _arg('dataset_base_path', None) + if enable_fov_retrieval: + fov_retriever = setup_fov_retriever_for_training( + dataset_base_path=dataset_base_path, + enable_fov_retrieval=True + ) + + launch_training_task( + dataset, model, model_logger, optimizer, scheduler, + num_epochs=args.num_epochs, + gradient_accumulation_steps=args.gradient_accumulation_steps, + per_device_train_batch_size=int(_arg("per_device_train_batch_size", 1)), + spike_threshold=_arg('spike_threshold', 5.0), + resume_step_count=resume_step_count, + enable_fov_retrieval=enable_fov_retrieval, + retrieval_method=_arg('retrieval_method', 'fov'), + latent_retrieval_dir=_arg('latent_retrieval_dir', None), + dataset_base_path=_arg('dataset_base_path', None), + fov_retriever=fov_retriever, + context_memory_frames=_arg('context_memory_frames', 8), + prev_chunk_frames=int(_arg('prev_chunk_frames', 81)), + fov_top_k=_arg('fov_top_k', 4), # Number of overlap frames (4), GT frame 0 added automatically + use_rt_relative=_arg('use_rt_relative', False), # Experiment 1_4_2: RT relative conversion + strict_overlap_context=_arg('strict_overlap_context', False), + dataset_repeat=_arg('dataset_repeat', 1), # Pass dataset_repeat for step calculation + use_camera_encoder=_arg('use_camera_encoder', False), # exp1_4_3: DDP find_unused_parameters + num_workers=_arg('num_workers', 0), + context_source=_arg('context_source', 'fov'), + max_train_steps=int(_arg('max_train_steps', 0)), + progress_total_steps=int(_arg('progress_total_steps', 0)), + log_interval=max(1, int(_arg('log_interval', 20) or 20)), + max_grad_norm=float(_arg('max_grad_norm', 1.0)), + ) + + model_logger.finish() \ No newline at end of file diff --git a/code/src/model_training/train.py.gradckpt-bak b/code/src/model_training/train.py.gradckpt-bak new file mode 100644 index 0000000000000000000000000000000000000000..81c4d72b6f8f5d8c77b8d96436789f2352f0f0b2 --- /dev/null +++ b/code/src/model_training/train.py.gradckpt-bak @@ -0,0 +1,660 @@ +import os, sys, re +import torch +import torch.nn as nn +import logging + +logger = logging.getLogger(__name__) + +_rank_env = os.environ.get("RANK") or os.environ.get("LOCAL_RANK") or os.environ.get("ACCELERATE_PROCESS_INDEX") or "0" +_rank = int(str(_rank_env)) +_level = logging.INFO if _rank == 0 else logging.WARNING +logging.basicConfig( + level=_level, + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + force=True, +) +logger.setLevel(_level) + +current_file_abs = os.path.abspath(__file__) +project_root = os.path.dirname(os.path.dirname(os.path.dirname(current_file_abs))) + +if project_root not in sys.path: + sys.path.insert(0, project_root) + +modules_to_clear = [ + 'diffsynth.models.memory.framepack_length', + 'diffsynth.models.memory.framepack_weight', + 'diffsynth.models.memory.spatial_grid_memory', + 'diffsynth.models.memory.videossm_hybrid', + 'diffsynth.models.memory.block_wise_ssm', + 'diffsynth.models.memory', + 'diffsynth.pipelines.wan_video_new', + 'diffsynth.trainers.utils', + 'diffsynth.models.wan_video_dit', + 'diffsynth.lora.flux_lora', + 'diffsynth.lora', + 'diffsynth.configs.model_config', + 'diffsynth.configs', + 'diffsynth.pipelines', + 'diffsynth.trainers', + 'diffsynth.models', + 'diffsynth', +] + +for mod in modules_to_clear: + if mod in sys.modules: + del sys.modules[mod] + +import importlib +importlib.invalidate_caches() + +from diffsynth.pipelines.wan_video_new import WanVideoPipeline, ModelConfig + +try: + import transformers + if not hasattr(transformers, "HybridCache") and hasattr(transformers, "DynamicCache"): + transformers.HybridCache = transformers.DynamicCache +except Exception: + pass + +from diffsynth.trainers.utils import DiffusionTrainingModule, ModelLogger as BaseModelLogger, VideoDataset, CamVideoDataset, wan_parser +from diffsynth.models.wan_video_dit import SelfAttention, CrossAttention, GateModule, modulate +from diffsynth.models.memory.videossm_hybrid import HybridStateSpaceMemory +from diffsynth.models.memory.block_wise_ssm import BlockWiseStateSpaceMemory +from diffsynth.models.memory.u_vit_cgla_blocks import CGLATransformerBlock, remap_wan_to_cgla + +try: + import diffsynth.trainers.utils as utils_module + utils_file = utils_module.__file__ if hasattr(utils_module, '__file__') else 'unknown' + is_local = 'site-packages' not in utils_file + if is_local: + logger.info(f"[VERIFIED] Using LOCAL diffsynth code from: {utils_file}") + else: + logger.warning(f"Using INSTALLED diffsynth package from: {utils_file}") +except Exception as e: + logger.error(f"Failed to verify code location: {e}") + +import random +import numpy as np +os.environ["TOKENIZERS_PARALLELISM"] = "false" +from safetensors.torch import load_file as safe_load_file +from src.model_training.fov_retrieval import setup_fov_retriever_for_training +from src.model_training.training_modules import DiTBlock_w_Action, WanTrainingModule + + +def set_seed(seed=42): + """Set random seeds for reproducible training.""" + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + os.environ['PYTHONHASHSEED'] = str(seed) + os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8' + logger.info(f"Random seed set to {seed}") + + +def _log_dit_freeze_summary(dit: torch.nn.Module) -> None: + by_module: dict[str, tuple[int, bool]] = {} + for name, p in dit.named_parameters(): + numel = p.numel() + trainable = p.requires_grad + parts = name.split(".") + prefix = ".".join(parts[:-1]) if len(parts) > 1 else name + if prefix not in by_module: + by_module[prefix] = (0, False) + prev_numel, prev_trainable = by_module[prefix] + by_module[prefix] = (prev_numel + numel, prev_trainable or trainable) + trainable_list = [(k, v[0]) for k, v in by_module.items() if v[1]] + frozen_list = [(k, v[0]) for k, v in by_module.items() if not v[1]] + trainable_list.sort(key=lambda x: x[0]) + frozen_list.sort(key=lambda x: x[0]) + total_trainable = sum(n for _, n in trainable_list) + total_frozen = sum(n for _, n in frozen_list) + examples = ", ".join(name for name, _ in trainable_list[:8]) + logger.info( + f"[DiT freeze] trainable={total_trainable:,} ({len(trainable_list)} groups), " + f"frozen={total_frozen:,} ({len(frozen_list)} groups), examples=[{examples}]" + ) + + +set_seed(42) + + +from src.model_training.training_modules.model_logger import ModelLogger +from src.model_training.training_modules.training_loop import launch_training_task + + +if __name__ == "__main__": + parser = wan_parser() + def _add_arg_if_missing(*args, **kwargs): + if args and args[0] in parser._option_string_actions: + return + parser.add_argument(*args, **kwargs) + + for name, kwargs in [ + ("--tokenizer_path", dict(type=str, default=None, help="Local tokenizer path.")), + ("--wandb_run_name", dict(type=str, default=None)), + ("--ckpt_interval", dict(type=int, default=None)), + ("--trainable_dit_modules", dict(type=str, default=None, help="Comma-separated DiT modules to unfreeze.")), + ("--num_workers", dict(type=int, default=0, help="DataLoader workers.")), + ("--max_train_steps", dict(type=int, default=0, help="Stop after N optimizer steps.")), + ("--progress_total_steps", dict(type=int, default=0, help="tqdm total steps override.")), + ("--log_interval", dict(type=int, default=20, help="Log loss/grad_norm/lr every N optimizer steps.")), + ("--max_grad_norm", dict(type=float, default=1.0, help="Grad-norm clip (default 1.0; 0 = no clip, norm still logged).")), + ("--resume_from_checkpoint", dict(type=str, default=None)), + ("--context_memory_frames", dict(type=int, default=8)), + ("--training_mode", dict(type=str, default="predict", choices=["predict", "context", "condition"])), + ("--context_drop_prob", dict(type=float, default=0.0)), + ("--retrieval_method", dict(type=str, default="fov", choices=["fov", "latent_sim"])), + ("--latent_retrieval_dir", dict(type=str, default=None)), + ("--fov_top_k", dict(type=int, default=4)), + ("--context_attention_weight", dict(type=float, default=1.0)), + ("--context_temporal_decay", dict(type=float, default=1.0)), + ("--spike_threshold", dict(type=float, default=5.0)), + ("--spatial_memory_tokens", dict(type=int, default=64)), + ("--spatial_memory_grid", dict(type=int, default=8)), + ("--spatial_memory_inject_mode", dict(type=str, default="concat_text", choices=["concat_text", "none", "cross_attn_readout"])), + ("--geometry_memory_column", dict(type=str, default="geometry_memory")), + ("--geometry_memory_root", dict(type=str, default=None)), + ("--geometry_spatial_memory_tokens", dict(type=int, default=64)), + ("--geometry_spatial_memory_grid", dict(type=int, default=8)), + ("--geometry_spatial_memory_temporal_bins", dict(type=int, default=4)), + ("--geometry_spatial_memory_inject_mode", dict(type=str, default="concat_text", choices=["concat_text", "none", "cross_attn_readout"])), + ("--framepack_ratio", dict(type=int, default=2)), + ("--framepack_length_strategy", dict(type=str, default="distance_merge", choices=["distance_merge", "mean", "uniform", "recent_weighted", "weighted_recent", "packed_multiscale"])), + ("--framepack_recent_keep_ratio", dict(type=float, default=0.5)), + ("--framepack_multiscale_w2", dict(type=float, default=0.25)), + ("--framepack_multiscale_w4", dict(type=float, default=0.15)), + ("--context_source", dict(type=str, default="fov", choices=["fov", "replay", "prev_chunk_tail"])), + ("--ssm_num_blocks_hint", dict(type=int, default=21)), + ("--ssm_every_n_blocks", dict(type=int, default=4)), + ("--videossm_kernel_size", dict(type=int, default=3)), + ("--videossm_expand", dict(type=int, default=2)), + ("--videossm_every_n_blocks", dict(type=int, default=4)), + # Camera-Guided Linear Attention (CGLA) memory row. + ("--cgla_every_n_blocks", dict(type=int, default=4)), + ("--cgla_num_heads", dict(type=int, default=0)), # 0 -> auto (dim // head_dim) + ("--cgla_head_dim", dict(type=int, default=128)), + ("--cgla_num_sparse_partition", dict(type=int, default=4)), + ("--cgla_num_writer", dict(type=int, default=1)), + ("--cgla_num_reader", dict(type=int, default=1)), + ("--cgla_gate_logit_normalizer", dict(type=int, default=16)), + ("--cgla_gate_low_rank_dim", dict(type=int, default=16)), + ("--cgla_pose_dim", dict(type=int, default=12)), + ("--cgla_pose_bottleneck", dict(type=int, default=64)), + ("--cgla_aux_loss_weight", dict(type=float, default=0.0)), + ("--cgla_mechanism", dict(type=str, default="cgla", + choices=["cgla", "prope", "ucpe"])), + ("--cgla_emb_dim", dict(type=int, default=1024)), + ("--cgla_num_patches", dict(type=int, default=880)), + ("--cgla_temporal_length", dict(type=int, default=21)), + ("--cgla_ffn_dim", dict(type=int, default=0)), # 0 -> use Wan DiT ffn_dim + ("--sampling_interval_steps", dict(type=int, default=0)), + ("--sampling_negative_prompt", dict(type=str, default="oversaturated colors, overexposed, static, blurry details")), + ("--sampling_height", dict(type=int, default=352)), + ("--sampling_width", dict(type=int, default=640)), + ("--sampling_num_frames", dict(type=int, default=81)), + ("--sampling_num_inference_steps", dict(type=int, default=50)), + ("--sampling_action_path", dict(type=str, default=None)), + ("--sampling_two_chunk_action_path", dict(type=str, default=None)), + ("--sampling_eval_dataset_base", dict(type=str, default=None)), + ("--sampling_eval_metadata_path", dict(type=str, default=None)), + ("--samples_per_epoch", dict(type=int, default=0)), + ("--camera_encoder_scale", dict(type=float, default=1.0)), + ("--camera_inject_mode", dict(type=str, default="post", choices=["post", "pre_norm", "pre_qkv", "pre_qkv_post", "pre_modulate", "pre_qkv_gated"])), + ]: + _add_arg_if_missing(name, **kwargs) + + for name in [ + "--save_full_model", "--add_action_attn", "--action_use_temporal_attention", + "--action_inject_after_spatial_attn", "--use_camera_encoder", "--camera_encoder_shallow", + "--camera_encoder_separate_t_r", "--camera_encoder_explicit_yaw", "--yaw_flip_aug", + "--camera_encoder_sincos_yaw", "--camera_encoder_r_mlp_no_layernorm", + "--add_camera_outside_gate", "--no_camera_encoder_zero_init", + "--camera_encoder_full_zero_init", "--enable_context_memory", "--context_per_frame_vae", + "--cfg_target_only", "--enable_fov_retrieval", "--use_rt_relative", + "--strict_overlap_context", "--use_anchor_frame", "--use_spatial_memory", + "--use_geometry_spatial_memory", + "--use_spatial_memory_legacy", "--use_framepack_memory", "--use_framepack_length_compress", + "--use_block_wise_ssm", "--use_videossm_hybrid", "--sampling_two_chunk_memory", + "--use_cgla_memory", "--cgla_use_pose_rope", "--cgla_use_pose_gate_mod", + ]: + _add_arg_if_missing(name, action="store_true") + + for name, kwargs in [ + ("--per_device_train_batch_size", dict(type=int, default=None)), + ("--timestep_shift", dict(type=float, default=1.0)), + ("--action_base_path", dict(type=str, default=None)), + ("--ckpt_path", dict(type=str, default=None)), + ("--cam_position_scale", dict(type=float, default=0.01)), + ("--resume_from", dict(type=str, default=None)), + ("--verify_ckpt_step", dict(type=int, default=0)), + ("--verify_high_noise_first_steps", dict(type=int, default=0)), + ("--moc_temperature", dict(type=float, default=1.0)), + ("--moc_top_k", dict(type=int, default=0)), + ("--prev_chunk_frames", dict(type=int, default=81)), + ("--implicit_type", dict(type=str, default="summary")), + ("--context_compressor_ratio", dict(type=int, default=2)), + ("--episodic_buffer_size", dict(type=int, default=0)), + ("--episodic_replay_interval", dict(type=int, default=0)), + ("--episodic_replay_weight", dict(type=float, default=0.0)), + ]: + _add_arg_if_missing(name, **kwargs) + for name in [ + "--enable_video_sampling", "--sampling_atomic_left_right", "--sampling_four_prompts", + "--sampling_two_prompts", "--train_action_module", "--train_cam_pose", + "--action_module_only", "--use_moc", "--unified_implicit", "--use_implicit_memory", + "--use_memory_v2v_compressor", "--use_slow_fast_memory", "--use_entity_memory", + "--use_episodic_memory", + ]: + _add_arg_if_missing(name, action="store_true") + args = parser.parse_args() + def _arg(name, default=None): + return getattr(args, name, default) + + def _normalize_and_validate_args(): + # Backward-compat mappings + if _arg("per_device_train_batch_size", None) is None: + args.per_device_train_batch_size = int(_arg("batch_size", 1)) + if _arg("sampling_atomic_left_right", False) and not _arg("sampling_two_chunk_memory", False): + # Legacy monitor intent maps to current two-chunk monitor. + args.sampling_two_chunk_memory = True + if _arg("enable_video_sampling", False) and int(_arg("sampling_interval_steps", 0)) <= 0: + args.sampling_interval_steps = 1000 + + # Keep paper-style block-wise SSM and legacy VideoSSM hybrid explicitly separated. + if _arg("use_block_wise_ssm", False) and _arg("use_videossm_hybrid", False): + raise ValueError( + "--use_block_wise_ssm and --use_videossm_hybrid are mutually exclusive; " + "use block-wise SSM for paper-aligned runs or VideoSSM hybrid for legacy baselines." + ) + + # CGLA is another per-block temporal-memory module on the same hook as + # block-wise SSM / VideoSSM hybrid; only one may be active at a time. + _per_block_memory_active = ( + bool(_arg("use_block_wise_ssm", False)) + or bool(_arg("use_videossm_hybrid", False)) + ) + if _arg("use_cgla_memory", False) and _per_block_memory_active: + raise ValueError( + "--use_cgla_memory is mutually exclusive with --use_block_wise_ssm / " + "--use_videossm_hybrid (two per-block temporal-memory modules on the same hook)." + ) + if _arg("use_cgla_memory", False) and not _arg("train_cam_pose", False): + logger.warning( + "--use_cgla_memory expects per-frame camera pose; enabling without " + "--train_cam_pose means no RT reaches the module (it will run as vanilla GLA)." + ) + + # Explicit retrieval strategy visibility: default fov, latent_sim degrades to fov when cache dir is absent. + if _arg("retrieval_method", "fov") == "latent_sim": + if not _arg("latent_retrieval_dir", None): + logger.warning("retrieval_method=latent_sim but latent_retrieval_dir is empty; runtime will fallback to fov retrieval.") + else: + logger.info(f"retrieval_method=latent_sim latent_retrieval_dir={args.latent_retrieval_dir}") + else: + logger.info("retrieval_method=fov") + + # 2-chunk sampling defaults: keep left/right_45 semantics compatible with existing shell wrappers. + if _arg("sampling_two_chunk_action_path", None) in (None, ""): + args.sampling_two_chunk_action_path = _arg("sampling_action_path", None) + + _normalize_and_validate_args() + + resume_step_count = 0 + if args.resume_from_checkpoint is not None: + if (_arg('trainable_dit_modules', None) or "").strip() or _arg('resume_weights_only', False): + logger.info("resume_from_checkpoint used for weights only (trainable_dit_modules set or resume_weights_only), step count starts from 0, no skip data") + resume_step_count = 0 + else: + checkpoint_filename = os.path.basename(args.resume_from_checkpoint) + step_match = re.search(r'Step-(\d+)', checkpoint_filename) + epoch_match = re.search(r'epoch-(\d+)', checkpoint_filename) + if step_match: + resume_step_count = int(step_match.group(1)) + logger.info(f"Resuming from step {resume_step_count} (extracted from checkpoint filename)") + elif epoch_match: + logger.info(f"Resuming from epoch checkpoint (epoch-{epoch_match.group(1)}), step count will start from 0") + resume_step_count = 0 + else: + logger.warning("Could not extract step count from checkpoint filename, starting from step 0") + + set_seed(42) + + args.enable_icl = False + args.icl_num_examples = 2 + args.icl_context_frames = 8 + + if _arg('train_cam_pose', False): + dataset = CamVideoDataset(args=args) + else: + dataset = VideoDataset(args=args, action_base_path=args.action_base_path) + + def _log_dataset_validation(ds): + ds_size = len(ds) + ds_repeat = _arg('dataset_repeat', 1) + logger.info( + f"[Dataset] size={ds_size}, repeat={ds_repeat}, " + f"epochs={args.num_epochs}, total_samples={ds_size * ds_repeat * args.num_epochs}" + ) + + _log_dataset_validation(dataset) + + model = WanTrainingModule( + model_paths=args.model_paths, + model_id_with_origin_paths=args.model_id_with_origin_paths, + tokenizer_path=_arg('tokenizer_path', None), + trainable_models=_arg('trainable_models', None), + lora_base_model=args.lora_base_model, + lora_target_modules=args.lora_target_modules, + lora_rank=args.lora_rank, + use_gradient_checkpointing_offload=args.use_gradient_checkpointing_offload, + extra_inputs=args.extra_inputs, + resume_from_checkpoint=args.resume_from_checkpoint, + dataset_base_path=_arg('dataset_base_path', None), + enable_context_memory=_arg('enable_context_memory', False), + context_drop_prob=_arg('context_drop_prob', 0.0), + context_drop_seed=42, + omit_context_actions=_arg('omit_context_actions', False) or (_arg('context_memory_frames', 8) == 1), # ctx=1: no context action injection + context_noise_prob=_arg('context_noise_prob', 0.0), + context_noise_std=_arg('context_noise_std', 0.02), + context_fixed_noise_std=_arg('context_fixed_noise_std', None), + context_memory_frames=_arg('context_memory_frames', 8), + context_per_frame_vae=_arg('context_per_frame_vae', False), + training_mode=_arg('training_mode', 'predict'), + teacher_forcing_prob=_arg('teacher_forcing_prob', 0.0), + yaw_flip_aug=_arg('yaw_flip_aug', False), + context_source=_arg('context_source', 'fov'), + use_framepack_memory=_arg('use_framepack_memory', False), + context_temporal_decay=_arg('context_temporal_decay', 1.0), + context_attention_weight=_arg('context_attention_weight', 1.0), + use_framepack_length_compress=_arg('use_framepack_length_compress', False), + framepack_ratio=_arg('framepack_ratio', 2), + framepack_length_strategy=_arg('framepack_length_strategy', 'distance_merge'), + framepack_recent_keep_ratio=_arg('framepack_recent_keep_ratio', 0.5), + framepack_multiscale_w2=_arg('framepack_multiscale_w2', 0.25), + framepack_multiscale_w4=_arg('framepack_multiscale_w4', 0.15), + use_spatial_memory=_arg('use_spatial_memory', False), + use_spatial_memory_legacy=_arg('use_spatial_memory_legacy', False), + spatial_memory_tokens=_arg('spatial_memory_tokens', 64), + spatial_memory_grid=_arg('spatial_memory_grid', 8), + spatial_memory_inject_mode=_arg('spatial_memory_inject_mode', 'concat_text'), + use_geometry_spatial_memory=_arg('use_geometry_spatial_memory', False), + geometry_spatial_memory_tokens=_arg('geometry_spatial_memory_tokens', 64), + geometry_spatial_memory_grid=_arg('geometry_spatial_memory_grid', 8), + geometry_spatial_memory_temporal_bins=_arg('geometry_spatial_memory_temporal_bins', 4), + geometry_spatial_memory_inject_mode=_arg( + 'geometry_spatial_memory_inject_mode', + 'concat_text', + ), + use_moc=_arg('use_moc', False), + moc_temperature=_arg('moc_temperature', 1.0), + moc_top_k=_arg('moc_top_k', 0), + timestep_shift=float(_arg('timestep_shift', 1.0)), + ) + if _arg('use_moc', False): + logger.info( + f"[MoC] enabled with temperature={float(_arg('moc_temperature', 1.0))}, " + f"top_k={int(_arg('moc_top_k', 0) or 0)}" + ) + if _arg('use_geometry_spatial_memory', False): + logger.info( + "[Geometry Spatial Memory] enabled; expects TSDF/point-cloud rendered condition " + f"from metadata column '{_arg('geometry_memory_column', 'geometry_memory')}'" + ) + + # ── VWM-style: Replace DiT blocks with DiTBlock_w_Action ── + _use_cam_pose = bool(_arg('train_cam_pose', False)) + use_cgla_memory = bool(_arg('use_cgla_memory', False)) + if _arg('train_action_module', False) or _use_cam_pose or use_cgla_memory: + dit = model.pipe.dit + old_blocks = dit.blocks + has_image_input = dit.has_image_input + dim = dit.dim + num_heads = dit.num_heads + ffn_dim = dit.ffn_dim + eps = 1e-6 + + block_dtype = next(old_blocks[0].parameters()).dtype + + use_block_wise_ssm = bool(_arg('use_block_wise_ssm', False)) + use_videossm_hybrid = bool(_arg('use_videossm_hybrid', False)) + ssm_every_n = max(int(_arg('ssm_every_n_blocks', 4)), 1) + videossm_every_n = max(int(_arg('videossm_every_n_blocks', 4)), 1) + # CGLA: which DiT blocks become CGLATransformerBlock (the rest stay + # DiTBlock_w_Action = Wan softmax + cam-pose). cgla_every_n_blocks=1 => + # all blocks (matches tests/test_cgla_wan.py); =4 => every 4th block, + # the same attach cadence as the SSM/VideoSSM rows (controlled ablation). + cgla_every_n = max(int(_arg('cgla_every_n_blocks', 4) or 4), 1) + _cgla_head_dim = int(_arg('cgla_head_dim', 128) or 128) + _cgla_head_dim = (_cgla_head_dim if _cgla_head_dim > 0 else (dim // num_heads)) + + new_blocks = nn.ModuleList() + for block_id, old_block in enumerate(old_blocks): + attach_block_ssm = use_block_wise_ssm and (block_id % ssm_every_n == 0) + attach_videossm = use_videossm_hybrid and (block_id % videossm_every_n == 0) + attach_cgla = use_cgla_memory and (block_id % cgla_every_n == 0) + if attach_cgla: + # CGLA row: the DiT block IS CGLATransformerBlock (Wan DiTBlock + # with self-attn swapped for the SSE-GLA). Wan's self_attn.{q,k,v,o} + # initialise the SSE-GLA {q,k,v,o}_proj (same shapes, via remap); + # cross_attn/norm/ffn/modulation load directly. + new_block = CGLATransformerBlock( + has_image_input=has_image_input, + dim=dim, num_heads=num_heads, ffn_dim=ffn_dim, eps=eps, + head_dim=_cgla_head_dim, + num_sparse_partition=int(_arg('cgla_num_sparse_partition', 4) or 4), + num_writer=int(_arg('cgla_num_writer', 1) or 1), + num_reader=int(_arg('cgla_num_reader', 1) or 1), + pose_dim=int(_arg('cgla_pose_dim', 12) or 12), + pose_bottleneck=int(_arg('cgla_pose_bottleneck', 64) or 64), + gate_logit_normalizer=int(_arg('cgla_gate_logit_normalizer', 16) or 16), + gate_low_rank_dim=int(_arg('cgla_gate_low_rank_dim', 16) or 16), + use_pose_gate_mod=bool(_arg('cgla_use_pose_gate_mod', False)), + layer_idx=block_id, + mechanism=str(_arg('cgla_mechanism', 'cgla') or 'cgla'), + bidirectional=True, + add_action_attn=_arg('add_action_attn', False), + action_use_temporal_attention=_arg('action_use_temporal_attention', False), + use_cam_pose=_use_cam_pose, + ) + new_block = new_block.to(dtype=block_dtype, device=next(old_block.parameters()).device) + # Copy the Wan submodules that match by name (cross_attn/norm/ffn). + for attr in ("cross_attn", "norm1", "norm2", "norm3", "ffn"): + if hasattr(old_block, attr) and hasattr(new_block, attr): + getattr(new_block, attr).load_state_dict(getattr(old_block, attr).state_dict()) + if hasattr(old_block, "modulation") and hasattr(new_block, "modulation"): + with torch.no_grad(): + new_block.modulation.copy_(old_block.modulation.to(dtype=block_dtype)) + # # self_attn: SSE-GLA structure differs from Wan SelfAttention, + # # so remap Wan's {q,k,v,o} -> {q,k,v,o}_proj and load (norm_q/ + # # norm_k have no SSE-GLA counterpart -> dropped). + # _sa_sd = {} + # for _k, _v in old_block.self_attn.state_dict().items(): + # if _k.startswith("q."): + # _sa_sd["q_proj." + _k[2:]] = _v + # elif _k.startswith("k."): + # _sa_sd["k_proj." + _k[2:]] = _v + # elif _k.startswith("v."): + # _sa_sd["v_proj." + _k[2:]] = _v + # elif _k.startswith("o."): + # _sa_sd["o_proj." + _k[2:]] = _v + # new_block.self_attn.load_state_dict(_sa_sd, strict=False) + else: + new_block = DiTBlock_w_Action( + has_image_input=has_image_input, + dim=dim, num_heads=num_heads, ffn_dim=ffn_dim, eps=eps, + add_action_attn=_arg('add_action_attn', False), + action_use_temporal_attention=_arg('action_use_temporal_attention', False), + use_cam_pose=_use_cam_pose, + use_block_wise_ssm=attach_block_ssm, + use_videossm_hybrid=attach_videossm, + videossm_kernel_size=int(_arg('videossm_kernel_size', 3) or 3), + videossm_expand=int(_arg('videossm_expand', 2) or 2), + ) + new_block = new_block.to(dtype=block_dtype, device=next(old_block.parameters()).device) + for attr in ("self_attn", "cross_attn", "norm1", "norm2", "norm3", "ffn"): + if hasattr(old_block, attr) and hasattr(new_block, attr): + getattr(new_block, attr).load_state_dict(getattr(old_block, attr).state_dict()) + if hasattr(old_block, "modulation") and hasattr(new_block, "modulation"): + with torch.no_grad(): + new_block.modulation.copy_(old_block.modulation.to(dtype=block_dtype)) + new_blocks.append(new_block) + + dit.blocks = new_blocks + _n_cgla = sum(1 for b in new_blocks if isinstance(b, CGLATransformerBlock)) + if use_cgla_memory: + _mech = str(_arg('cgla_mechanism', 'cgla') or 'cgla') + logger.info(f"[CGLA] Replaced {_n_cgla}/{len(new_blocks)} DiT blocks with " + f"CGLATransformerBlock (every_n={cgla_every_n}, mechanism={_mech}, " + f"head_dim={_cgla_head_dim}, bidirectional); " + f"Wan self_attn q/k/v/o -> SSE-GLA q/k/v/o_proj") + + _mlp_type = "MLP_CamPose" if _use_cam_pose else "MLP_Action" + logger.info(f"[VWM-style] Replaced {len(new_blocks)} DiT blocks with DiTBlock_w_Action ({_mlp_type}, zero-init)") + if use_block_wise_ssm: + logger.info(f"[Block-wise SSM] attached to every {ssm_every_n} DiT block(s)") + if use_videossm_hybrid: + logger.info(f"[VideoSSM hybrid] attached to every {videossm_every_n} DiT block(s)") + + device = next(dit.parameters()).device + _ckpt_path = _arg('ckpt_path', None) or _arg('resume_from_checkpoint', None) + if _ckpt_path is not None: + ckpt = safe_load_file(_ckpt_path) + if use_cgla_memory: + # Remap Wan self_attn.{q,k,v,o} -> SSE-GLA {q,k,v,o}_proj so a Wan + # checkpoint loads the linear-attention projections. No-op on an + # already-CGLA checkpoint (keys are already *_proj). + ckpt = remap_wan_to_cgla(ckpt) + missing, unexpected = dit.load_state_dict(ckpt, strict=False) + dit.to(device=device) + logger.info(f"[VWM-style] Loaded ckpt: {len(ckpt)} keys, missing={len(missing)}, unexpected={len(unexpected)}") + + if use_cgla_memory: + # Block-aware freeze (supports cgla_every_n_blocks > 1, where some + # blocks stay DiTBlock_w_Action = Wan softmax). CGLA blocks: train + # the SSE-GLA self_attn + cam_encoder + cgla_gate + noise_write_gate, + # freeze the Wan backbone (cross_attn/norm/ffn/modulation). Non-CGLA + # blocks: VWM pattern (train action_mlp / self_attn_with_action / + # SSM, freeze Wan softmax self_attn + backbone). + for block in dit.blocks: + if isinstance(block, CGLATransformerBlock): + for name, param in block.named_parameters(): + if ("action_mlp" in name) or ("self_attn_with_action" in name) or ("self_attn" in name) or \ + ("cam_encoder" in name) or ("noise_write_gate" in name): + param.requires_grad = True + else: + param.requires_grad = False + else: + for name, param in block.named_parameters(): + if ("action_mlp" in name) or ("self_attn_with_action" in name) or \ + ("block_wise_ssm" in name) or ("videossm_hybrid" in name): + param.requires_grad = True + else: + param.requires_grad = False + elif _arg('action_module_only', False): + if _arg('add_action_attn', False): + for block in dit.blocks: + for name, param in block.named_parameters(): + if ("action_mlp" in name) or ("self_attn_with_action" in name) or ("block_wise_ssm" in name) or ("videossm_hybrid" in name): + param.requires_grad = True + else: + param.requires_grad = False + else: + for block in dit.blocks: + for name, param in block.named_parameters(): + if "action_mlp" in name or "self_attn" in name or "block_wise_ssm" in name or "videossm_hybrid" in name: + param.requires_grad = True + else: + param.requires_grad = False + else: + for block in dit.blocks: + for name, param in block.named_parameters(): + if "action_mlp" in name or "self_attn_with_action" in name or "block_wise_ssm" in name or "videossm_hybrid" in name: + param.requires_grad = True + else: + param.requires_grad = False + _log_dit_freeze_summary(dit) + + _resume_from = _arg('resume_from', None) + if _resume_from: + logger.info(f"Loading full resume checkpoint: {_resume_from}") + ckpt = safe_load_file(_resume_from) + if use_cgla_memory: + # No-op on an already-CGLA ckpt (keys are *_proj); remaps a Wan ckpt. + ckpt = remap_wan_to_cgla(ckpt) + model.pipe.dit.load_state_dict(ckpt, strict=False) + logger.info(f"Checkpoint loaded, resuming from step {resume_step_count}") + + model_logger = ModelLogger( + args.output_path, + remove_prefix_in_ckpt=args.remove_prefix_in_ckpt, + wandb_run_name=args.wandb_run_name, + ckpt_interval=args.ckpt_interval, + resume_step_count=resume_step_count, + save_full_model=_arg('save_full_model', False), + context_drop_prob=float(_arg("context_drop_prob", 0.0)), + enable_video_sampling=_arg("enable_video_sampling", False), + sampling_interval_steps=int(_arg("sampling_interval_steps", 0)), + sampling_two_chunk_memory=_arg("sampling_two_chunk_memory", False), + sampling_action_path=_arg("sampling_action_path", None), + sampling_two_chunk_action_path=_arg("sampling_two_chunk_action_path", None), + sampling_negative_prompt=_arg("sampling_negative_prompt", ""), + sampling_height=int(_arg("sampling_height", 352)), + sampling_width=int(_arg("sampling_width", 640)), + sampling_num_frames=int(_arg("sampling_num_frames", 81)), + sampling_num_inference_steps=int(_arg("sampling_num_inference_steps", 50)), + context_memory_frames=int(_arg("context_memory_frames", 1)), + context_source=_arg("context_source", "replay"), + context_per_frame_vae=_arg("context_per_frame_vae", False), + # Monitor samples with the SAME noise-schedule shift training uses + # (--timestep_shift), so the in-training video reflects the trained + # schedule (15 for two-chunk rows, 5 for legacy ctx rows). + sampling_sigma_shift=float(_arg("timestep_shift", 1.0) or 1.0), + ) + + optimizer = torch.optim.AdamW(model.trainable_modules(), lr=args.learning_rate) + scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer) + + # Setup FOV retriever for context-based memory training (also for ModelLogger sampling) + enable_fov_retrieval = _arg('enable_fov_retrieval', False) + fov_retriever = None + dataset_base_path = _arg('dataset_base_path', None) + if enable_fov_retrieval: + fov_retriever = setup_fov_retriever_for_training( + dataset_base_path=dataset_base_path, + enable_fov_retrieval=True + ) + + launch_training_task( + dataset, model, model_logger, optimizer, scheduler, + num_epochs=args.num_epochs, + gradient_accumulation_steps=args.gradient_accumulation_steps, + per_device_train_batch_size=int(_arg("per_device_train_batch_size", 1)), + spike_threshold=_arg('spike_threshold', 5.0), + resume_step_count=resume_step_count, + enable_fov_retrieval=enable_fov_retrieval, + retrieval_method=_arg('retrieval_method', 'fov'), + latent_retrieval_dir=_arg('latent_retrieval_dir', None), + dataset_base_path=_arg('dataset_base_path', None), + fov_retriever=fov_retriever, + context_memory_frames=_arg('context_memory_frames', 8), + prev_chunk_frames=int(_arg('prev_chunk_frames', 81)), + fov_top_k=_arg('fov_top_k', 4), # Number of overlap frames (4), GT frame 0 added automatically + use_rt_relative=_arg('use_rt_relative', False), # Experiment 1_4_2: RT relative conversion + strict_overlap_context=_arg('strict_overlap_context', False), + dataset_repeat=_arg('dataset_repeat', 1), # Pass dataset_repeat for step calculation + use_camera_encoder=_arg('use_camera_encoder', False), # exp1_4_3: DDP find_unused_parameters + num_workers=_arg('num_workers', 0), + context_source=_arg('context_source', 'fov'), + max_train_steps=int(_arg('max_train_steps', 0)), + progress_total_steps=int(_arg('progress_total_steps', 0)), + log_interval=max(1, int(_arg('log_interval', 20) or 20)), + max_grad_norm=float(_arg('max_grad_norm', 1.0)), + ) + + model_logger.finish() \ No newline at end of file diff --git a/code/src/model_training/train.py.p0bak b/code/src/model_training/train.py.p0bak new file mode 100644 index 0000000000000000000000000000000000000000..61cbddc0dad6a05ae3902f354992627b5cd0e62c --- /dev/null +++ b/code/src/model_training/train.py.p0bak @@ -0,0 +1,660 @@ +import os, sys, re +import torch +import torch.nn as nn +import logging + +logger = logging.getLogger(__name__) + +_rank_env = os.environ.get("RANK") or os.environ.get("LOCAL_RANK") or os.environ.get("ACCELERATE_PROCESS_INDEX") or "0" +_rank = int(str(_rank_env)) +_level = logging.INFO if _rank == 0 else logging.WARNING +logging.basicConfig( + level=_level, + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + force=True, +) +logger.setLevel(_level) + +current_file_abs = os.path.abspath(__file__) +project_root = os.path.dirname(os.path.dirname(os.path.dirname(current_file_abs))) + +if project_root not in sys.path: + sys.path.insert(0, project_root) + +modules_to_clear = [ + 'diffsynth.models.memory.framepack_length', + 'diffsynth.models.memory.framepack_weight', + 'diffsynth.models.memory.spatial_grid_memory', + 'diffsynth.models.memory.videossm_hybrid', + 'diffsynth.models.memory.block_wise_ssm', + 'diffsynth.models.memory', + 'diffsynth.pipelines.wan_video_new', + 'diffsynth.trainers.utils', + 'diffsynth.models.wan_video_dit', + 'diffsynth.lora.flux_lora', + 'diffsynth.lora', + 'diffsynth.configs.model_config', + 'diffsynth.configs', + 'diffsynth.pipelines', + 'diffsynth.trainers', + 'diffsynth.models', + 'diffsynth', +] + +for mod in modules_to_clear: + if mod in sys.modules: + del sys.modules[mod] + +import importlib +importlib.invalidate_caches() + +from diffsynth.pipelines.wan_video_new import WanVideoPipeline, ModelConfig + +try: + import transformers + if not hasattr(transformers, "HybridCache") and hasattr(transformers, "DynamicCache"): + transformers.HybridCache = transformers.DynamicCache +except Exception: + pass + +from diffsynth.trainers.utils import DiffusionTrainingModule, ModelLogger as BaseModelLogger, VideoDataset, CamVideoDataset, wan_parser +from diffsynth.models.wan_video_dit import SelfAttention, CrossAttention, GateModule, modulate +from diffsynth.models.memory.videossm_hybrid import HybridStateSpaceMemory +from diffsynth.models.memory.block_wise_ssm import BlockWiseStateSpaceMemory +from diffsynth.models.memory.u_vit_cgla_blocks import CGLATransformerBlock, remap_wan_to_cgla + +try: + import diffsynth.trainers.utils as utils_module + utils_file = utils_module.__file__ if hasattr(utils_module, '__file__') else 'unknown' + is_local = 'site-packages' not in utils_file + if is_local: + logger.info(f"[VERIFIED] Using LOCAL diffsynth code from: {utils_file}") + else: + logger.warning(f"Using INSTALLED diffsynth package from: {utils_file}") +except Exception as e: + logger.error(f"Failed to verify code location: {e}") + +import random +import numpy as np +os.environ["TOKENIZERS_PARALLELISM"] = "false" +from safetensors.torch import load_file as safe_load_file +from src.model_training.fov_retrieval import setup_fov_retriever_for_training +from src.model_training.training_modules import DiTBlock_w_Action, WanTrainingModule + + +def set_seed(seed=42): + """Set random seeds for reproducible training.""" + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + os.environ['PYTHONHASHSEED'] = str(seed) + os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8' + logger.info(f"Random seed set to {seed}") + + +def _log_dit_freeze_summary(dit: torch.nn.Module) -> None: + by_module: dict[str, tuple[int, bool]] = {} + for name, p in dit.named_parameters(): + numel = p.numel() + trainable = p.requires_grad + parts = name.split(".") + prefix = ".".join(parts[:-1]) if len(parts) > 1 else name + if prefix not in by_module: + by_module[prefix] = (0, False) + prev_numel, prev_trainable = by_module[prefix] + by_module[prefix] = (prev_numel + numel, prev_trainable or trainable) + trainable_list = [(k, v[0]) for k, v in by_module.items() if v[1]] + frozen_list = [(k, v[0]) for k, v in by_module.items() if not v[1]] + trainable_list.sort(key=lambda x: x[0]) + frozen_list.sort(key=lambda x: x[0]) + total_trainable = sum(n for _, n in trainable_list) + total_frozen = sum(n for _, n in frozen_list) + examples = ", ".join(name for name, _ in trainable_list[:8]) + logger.info( + f"[DiT freeze] trainable={total_trainable:,} ({len(trainable_list)} groups), " + f"frozen={total_frozen:,} ({len(frozen_list)} groups), examples=[{examples}]" + ) + + +set_seed(42) + + +from src.model_training.training_modules.model_logger import ModelLogger +from src.model_training.training_modules.training_loop import launch_training_task + + +if __name__ == "__main__": + parser = wan_parser() + def _add_arg_if_missing(*args, **kwargs): + if args and args[0] in parser._option_string_actions: + return + parser.add_argument(*args, **kwargs) + + for name, kwargs in [ + ("--tokenizer_path", dict(type=str, default=None, help="Local tokenizer path.")), + ("--wandb_run_name", dict(type=str, default=None)), + ("--ckpt_interval", dict(type=int, default=None)), + ("--trainable_dit_modules", dict(type=str, default=None, help="Comma-separated DiT modules to unfreeze.")), + ("--num_workers", dict(type=int, default=0, help="DataLoader workers.")), + ("--max_train_steps", dict(type=int, default=0, help="Stop after N optimizer steps.")), + ("--progress_total_steps", dict(type=int, default=0, help="tqdm total steps override.")), + ("--log_interval", dict(type=int, default=20, help="Log loss/grad_norm/lr every N optimizer steps.")), + ("--max_grad_norm", dict(type=float, default=1.0, help="Grad-norm clip (default 1.0; 0 = no clip, norm still logged).")), + ("--resume_from_checkpoint", dict(type=str, default=None)), + ("--context_memory_frames", dict(type=int, default=8)), + ("--training_mode", dict(type=str, default="predict", choices=["predict", "context", "condition"])), + ("--context_drop_prob", dict(type=float, default=0.0)), + ("--retrieval_method", dict(type=str, default="fov", choices=["fov", "latent_sim"])), + ("--latent_retrieval_dir", dict(type=str, default=None)), + ("--fov_top_k", dict(type=int, default=4)), + ("--context_attention_weight", dict(type=float, default=1.0)), + ("--context_temporal_decay", dict(type=float, default=1.0)), + ("--spike_threshold", dict(type=float, default=5.0)), + ("--spatial_memory_tokens", dict(type=int, default=64)), + ("--spatial_memory_grid", dict(type=int, default=8)), + ("--spatial_memory_inject_mode", dict(type=str, default="concat_text", choices=["concat_text", "none", "cross_attn_readout"])), + ("--geometry_memory_column", dict(type=str, default="geometry_memory")), + ("--geometry_memory_root", dict(type=str, default=None)), + ("--geometry_spatial_memory_tokens", dict(type=int, default=64)), + ("--geometry_spatial_memory_grid", dict(type=int, default=8)), + ("--geometry_spatial_memory_temporal_bins", dict(type=int, default=4)), + ("--geometry_spatial_memory_inject_mode", dict(type=str, default="concat_text", choices=["concat_text", "none", "cross_attn_readout"])), + ("--framepack_ratio", dict(type=int, default=2)), + ("--framepack_length_strategy", dict(type=str, default="distance_merge", choices=["distance_merge", "mean", "uniform", "recent_weighted", "weighted_recent", "packed_multiscale"])), + ("--framepack_recent_keep_ratio", dict(type=float, default=0.5)), + ("--framepack_multiscale_w2", dict(type=float, default=0.25)), + ("--framepack_multiscale_w4", dict(type=float, default=0.15)), + ("--context_source", dict(type=str, default="fov", choices=["fov", "replay", "prev_chunk_tail"])), + ("--ssm_num_blocks_hint", dict(type=int, default=21)), + ("--ssm_every_n_blocks", dict(type=int, default=4)), + ("--videossm_kernel_size", dict(type=int, default=3)), + ("--videossm_expand", dict(type=int, default=2)), + ("--videossm_every_n_blocks", dict(type=int, default=4)), + # Camera-Guided Linear Attention (CGLA) memory row. + ("--cgla_every_n_blocks", dict(type=int, default=4)), + ("--cgla_num_heads", dict(type=int, default=0)), # 0 -> auto (dim // head_dim) + ("--cgla_head_dim", dict(type=int, default=128)), + ("--cgla_num_sparse_partition", dict(type=int, default=4)), + ("--cgla_num_writer", dict(type=int, default=1)), + ("--cgla_num_reader", dict(type=int, default=1)), + ("--cgla_gate_logit_normalizer", dict(type=int, default=16)), + ("--cgla_gate_low_rank_dim", dict(type=int, default=16)), + ("--cgla_pose_dim", dict(type=int, default=12)), + ("--cgla_pose_bottleneck", dict(type=int, default=64)), + ("--cgla_aux_loss_weight", dict(type=float, default=0.0)), + ("--cgla_mechanism", dict(type=str, default="cgla", + choices=["cgla", "prope", "ucpe"])), + ("--cgla_emb_dim", dict(type=int, default=1024)), + ("--cgla_num_patches", dict(type=int, default=880)), + ("--cgla_temporal_length", dict(type=int, default=21)), + ("--cgla_ffn_dim", dict(type=int, default=0)), # 0 -> use Wan DiT ffn_dim + ("--sampling_interval_steps", dict(type=int, default=0)), + ("--sampling_negative_prompt", dict(type=str, default="oversaturated colors, overexposed, static, blurry details")), + ("--sampling_height", dict(type=int, default=352)), + ("--sampling_width", dict(type=int, default=640)), + ("--sampling_num_frames", dict(type=int, default=81)), + ("--sampling_num_inference_steps", dict(type=int, default=50)), + ("--sampling_action_path", dict(type=str, default=None)), + ("--sampling_two_chunk_action_path", dict(type=str, default=None)), + ("--sampling_eval_dataset_base", dict(type=str, default=None)), + ("--sampling_eval_metadata_path", dict(type=str, default=None)), + ("--samples_per_epoch", dict(type=int, default=0)), + ("--camera_encoder_scale", dict(type=float, default=1.0)), + ("--camera_inject_mode", dict(type=str, default="post", choices=["post", "pre_norm", "pre_qkv", "pre_qkv_post", "pre_modulate", "pre_qkv_gated"])), + ]: + _add_arg_if_missing(name, **kwargs) + + for name in [ + "--save_full_model", "--add_action_attn", "--action_use_temporal_attention", + "--action_inject_after_spatial_attn", "--use_camera_encoder", "--camera_encoder_shallow", + "--camera_encoder_separate_t_r", "--camera_encoder_explicit_yaw", "--yaw_flip_aug", + "--camera_encoder_sincos_yaw", "--camera_encoder_r_mlp_no_layernorm", + "--add_camera_outside_gate", "--no_camera_encoder_zero_init", + "--camera_encoder_full_zero_init", "--enable_context_memory", "--context_per_frame_vae", + "--cfg_target_only", "--enable_fov_retrieval", "--use_rt_relative", + "--strict_overlap_context", "--use_anchor_frame", "--use_spatial_memory", + "--use_geometry_spatial_memory", + "--use_spatial_memory_legacy", "--use_framepack_memory", "--use_framepack_length_compress", + "--use_block_wise_ssm", "--use_videossm_hybrid", "--sampling_two_chunk_memory", + "--use_cgla_memory", "--cgla_use_pose_rope", "--cgla_use_pose_gate_mod", + ]: + _add_arg_if_missing(name, action="store_true") + + for name, kwargs in [ + ("--per_device_train_batch_size", dict(type=int, default=None)), + ("--timestep_shift", dict(type=float, default=1.0)), + ("--action_base_path", dict(type=str, default=None)), + ("--ckpt_path", dict(type=str, default=None)), + ("--cam_position_scale", dict(type=float, default=0.01)), + ("--resume_from", dict(type=str, default=None)), + ("--verify_ckpt_step", dict(type=int, default=0)), + ("--verify_high_noise_first_steps", dict(type=int, default=0)), + ("--moc_temperature", dict(type=float, default=1.0)), + ("--moc_top_k", dict(type=int, default=0)), + ("--prev_chunk_frames", dict(type=int, default=81)), + ("--implicit_type", dict(type=str, default="summary")), + ("--context_compressor_ratio", dict(type=int, default=2)), + ("--episodic_buffer_size", dict(type=int, default=0)), + ("--episodic_replay_interval", dict(type=int, default=0)), + ("--episodic_replay_weight", dict(type=float, default=0.0)), + ]: + _add_arg_if_missing(name, **kwargs) + for name in [ + "--enable_video_sampling", "--sampling_atomic_left_right", "--sampling_four_prompts", + "--sampling_two_prompts", "--train_action_module", "--train_cam_pose", + "--action_module_only", "--use_moc", "--unified_implicit", "--use_implicit_memory", + "--use_memory_v2v_compressor", "--use_slow_fast_memory", "--use_entity_memory", + "--use_episodic_memory", + ]: + _add_arg_if_missing(name, action="store_true") + args = parser.parse_args() + def _arg(name, default=None): + return getattr(args, name, default) + + def _normalize_and_validate_args(): + # Backward-compat mappings + if _arg("per_device_train_batch_size", None) is None: + args.per_device_train_batch_size = int(_arg("batch_size", 1)) + if _arg("sampling_atomic_left_right", False) and not _arg("sampling_two_chunk_memory", False): + # Legacy monitor intent maps to current two-chunk monitor. + args.sampling_two_chunk_memory = True + if _arg("enable_video_sampling", False) and int(_arg("sampling_interval_steps", 0)) <= 0: + args.sampling_interval_steps = 1000 + + # Keep paper-style block-wise SSM and legacy VideoSSM hybrid explicitly separated. + if _arg("use_block_wise_ssm", False) and _arg("use_videossm_hybrid", False): + raise ValueError( + "--use_block_wise_ssm and --use_videossm_hybrid are mutually exclusive; " + "use block-wise SSM for paper-aligned runs or VideoSSM hybrid for legacy baselines." + ) + + # CGLA is another per-block temporal-memory module on the same hook as + # block-wise SSM / VideoSSM hybrid; only one may be active at a time. + _per_block_memory_active = ( + bool(_arg("use_block_wise_ssm", False)) + or bool(_arg("use_videossm_hybrid", False)) + ) + if _arg("use_cgla_memory", False) and _per_block_memory_active: + raise ValueError( + "--use_cgla_memory is mutually exclusive with --use_block_wise_ssm / " + "--use_videossm_hybrid (two per-block temporal-memory modules on the same hook)." + ) + if _arg("use_cgla_memory", False) and not _arg("train_cam_pose", False): + logger.warning( + "--use_cgla_memory expects per-frame camera pose; enabling without " + "--train_cam_pose means no RT reaches the module (it will run as vanilla GLA)." + ) + + # Explicit retrieval strategy visibility: default fov, latent_sim degrades to fov when cache dir is absent. + if _arg("retrieval_method", "fov") == "latent_sim": + if not _arg("latent_retrieval_dir", None): + logger.warning("retrieval_method=latent_sim but latent_retrieval_dir is empty; runtime will fallback to fov retrieval.") + else: + logger.info(f"retrieval_method=latent_sim latent_retrieval_dir={args.latent_retrieval_dir}") + else: + logger.info("retrieval_method=fov") + + # 2-chunk sampling defaults: keep left/right_45 semantics compatible with existing shell wrappers. + if _arg("sampling_two_chunk_action_path", None) in (None, ""): + args.sampling_two_chunk_action_path = _arg("sampling_action_path", None) + + _normalize_and_validate_args() + + resume_step_count = 0 + if args.resume_from_checkpoint is not None: + if (_arg('trainable_dit_modules', None) or "").strip() or _arg('resume_weights_only', False): + logger.info("resume_from_checkpoint used for weights only (trainable_dit_modules set or resume_weights_only), step count starts from 0, no skip data") + resume_step_count = 0 + else: + checkpoint_filename = os.path.basename(args.resume_from_checkpoint) + step_match = re.search(r'Step-(\d+)', checkpoint_filename) + epoch_match = re.search(r'epoch-(\d+)', checkpoint_filename) + if step_match: + resume_step_count = int(step_match.group(1)) + logger.info(f"Resuming from step {resume_step_count} (extracted from checkpoint filename)") + elif epoch_match: + logger.info(f"Resuming from epoch checkpoint (epoch-{epoch_match.group(1)}), step count will start from 0") + resume_step_count = 0 + else: + logger.warning("Could not extract step count from checkpoint filename, starting from step 0") + + set_seed(42) + + args.enable_icl = False + args.icl_num_examples = 2 + args.icl_context_frames = 8 + + if _arg('train_cam_pose', False): + dataset = CamVideoDataset(args=args) + else: + dataset = VideoDataset(args=args, action_base_path=args.action_base_path) + + def _log_dataset_validation(ds): + ds_size = len(ds) + ds_repeat = _arg('dataset_repeat', 1) + logger.info( + f"[Dataset] size={ds_size}, repeat={ds_repeat}, " + f"epochs={args.num_epochs}, total_samples={ds_size * ds_repeat * args.num_epochs}" + ) + + _log_dataset_validation(dataset) + + model = WanTrainingModule( + model_paths=args.model_paths, + model_id_with_origin_paths=args.model_id_with_origin_paths, + tokenizer_path=_arg('tokenizer_path', None), + trainable_models=_arg('trainable_models', None), + lora_base_model=args.lora_base_model, + lora_target_modules=args.lora_target_modules, + lora_rank=args.lora_rank, + use_gradient_checkpointing_offload=args.use_gradient_checkpointing_offload, + extra_inputs=args.extra_inputs, + resume_from_checkpoint=args.resume_from_checkpoint, + dataset_base_path=_arg('dataset_base_path', None), + enable_context_memory=_arg('enable_context_memory', False), + context_drop_prob=_arg('context_drop_prob', 0.0), + context_drop_seed=42, + omit_context_actions=_arg('omit_context_actions', False) or (_arg('context_memory_frames', 8) == 1), # ctx=1: no context action injection + context_noise_prob=_arg('context_noise_prob', 0.0), + context_noise_std=_arg('context_noise_std', 0.02), + context_fixed_noise_std=_arg('context_fixed_noise_std', None), + context_memory_frames=_arg('context_memory_frames', 8), + context_per_frame_vae=_arg('context_per_frame_vae', False), + training_mode=_arg('training_mode', 'predict'), + teacher_forcing_prob=_arg('teacher_forcing_prob', 0.0), + yaw_flip_aug=_arg('yaw_flip_aug', False), + context_source=_arg('context_source', 'fov'), + use_framepack_memory=_arg('use_framepack_memory', False), + context_temporal_decay=_arg('context_temporal_decay', 1.0), + context_attention_weight=_arg('context_attention_weight', 1.0), + use_framepack_length_compress=_arg('use_framepack_length_compress', False), + framepack_ratio=_arg('framepack_ratio', 2), + framepack_length_strategy=_arg('framepack_length_strategy', 'distance_merge'), + framepack_recent_keep_ratio=_arg('framepack_recent_keep_ratio', 0.5), + framepack_multiscale_w2=_arg('framepack_multiscale_w2', 0.25), + framepack_multiscale_w4=_arg('framepack_multiscale_w4', 0.15), + use_spatial_memory=_arg('use_spatial_memory', False), + use_spatial_memory_legacy=_arg('use_spatial_memory_legacy', False), + spatial_memory_tokens=_arg('spatial_memory_tokens', 64), + spatial_memory_grid=_arg('spatial_memory_grid', 8), + spatial_memory_inject_mode=_arg('spatial_memory_inject_mode', 'concat_text'), + use_geometry_spatial_memory=_arg('use_geometry_spatial_memory', False), + geometry_spatial_memory_tokens=_arg('geometry_spatial_memory_tokens', 64), + geometry_spatial_memory_grid=_arg('geometry_spatial_memory_grid', 8), + geometry_spatial_memory_temporal_bins=_arg('geometry_spatial_memory_temporal_bins', 4), + geometry_spatial_memory_inject_mode=_arg( + 'geometry_spatial_memory_inject_mode', + 'concat_text', + ), + use_moc=_arg('use_moc', False), + moc_temperature=_arg('moc_temperature', 1.0), + moc_top_k=_arg('moc_top_k', 0), + timestep_shift=float(_arg('timestep_shift', 1.0)), + ) + if _arg('use_moc', False): + logger.info( + f"[MoC] enabled with temperature={float(_arg('moc_temperature', 1.0))}, " + f"top_k={int(_arg('moc_top_k', 0) or 0)}" + ) + if _arg('use_geometry_spatial_memory', False): + logger.info( + "[Geometry Spatial Memory] enabled; expects TSDF/point-cloud rendered condition " + f"from metadata column '{_arg('geometry_memory_column', 'geometry_memory')}'" + ) + + # ── VWM-style: Replace DiT blocks with DiTBlock_w_Action ── + _use_cam_pose = bool(_arg('train_cam_pose', False)) + use_cgla_memory = bool(_arg('use_cgla_memory', False)) + if _arg('train_action_module', False) or _use_cam_pose or use_cgla_memory: + dit = model.pipe.dit + old_blocks = dit.blocks + has_image_input = dit.has_image_input + dim = dit.dim + num_heads = dit.num_heads + ffn_dim = dit.ffn_dim + eps = 1e-6 + + block_dtype = next(old_blocks[0].parameters()).dtype + + use_block_wise_ssm = bool(_arg('use_block_wise_ssm', False)) + use_videossm_hybrid = bool(_arg('use_videossm_hybrid', False)) + ssm_every_n = max(int(_arg('ssm_every_n_blocks', 4)), 1) + videossm_every_n = max(int(_arg('videossm_every_n_blocks', 4)), 1) + # CGLA: which DiT blocks become CGLATransformerBlock (the rest stay + # DiTBlock_w_Action = Wan softmax + cam-pose). cgla_every_n_blocks=1 => + # all blocks (matches tests/test_cgla_wan.py); =4 => every 4th block, + # the same attach cadence as the SSM/VideoSSM rows (controlled ablation). + cgla_every_n = max(int(_arg('cgla_every_n_blocks', 4) or 4), 1) + _cgla_head_dim = int(_arg('cgla_head_dim', 128) or 128) + _cgla_head_dim = (_cgla_head_dim if _cgla_head_dim > 0 else (dim // num_heads)) + + new_blocks = nn.ModuleList() + for block_id, old_block in enumerate(old_blocks): + attach_block_ssm = use_block_wise_ssm and (block_id % ssm_every_n == 0) + attach_videossm = use_videossm_hybrid and (block_id % videossm_every_n == 0) + attach_cgla = use_cgla_memory and (block_id % cgla_every_n == 0) + if attach_cgla: + # CGLA row: the DiT block IS CGLATransformerBlock (Wan DiTBlock + # with self-attn swapped for the SSE-GLA). Wan's self_attn.{q,k,v,o} + # initialise the SSE-GLA {q,k,v,o}_proj (same shapes, via remap); + # cross_attn/norm/ffn/modulation load directly. + new_block = CGLATransformerBlock( + has_image_input=has_image_input, + dim=dim, num_heads=num_heads, ffn_dim=ffn_dim, eps=eps, + head_dim=_cgla_head_dim, + num_sparse_partition=int(_arg('cgla_num_sparse_partition', 4) or 4), + num_writer=int(_arg('cgla_num_writer', 1) or 1), + num_reader=int(_arg('cgla_num_reader', 1) or 1), + pose_dim=int(_arg('cgla_pose_dim', 12) or 12), + pose_bottleneck=int(_arg('cgla_pose_bottleneck', 64) or 64), + gate_logit_normalizer=int(_arg('cgla_gate_logit_normalizer', 16) or 16), + gate_low_rank_dim=int(_arg('cgla_gate_low_rank_dim', 16) or 16), + use_pose_gate_mod=bool(_arg('cgla_use_pose_gate_mod', False)), + layer_idx=block_id, + mechanism=str(_arg('cgla_mechanism', 'cgla') or 'cgla'), + bidirectional=True, + add_action_attn=_arg('add_action_attn', False), + action_use_temporal_attention=_arg('action_use_temporal_attention', False), + use_cam_pose=_use_cam_pose, + ) + new_block = new_block.to(dtype=block_dtype, device=next(old_block.parameters()).device) + # Copy the Wan submodules that match by name (cross_attn/norm/ffn). + for attr in ("cross_attn", "norm1", "norm2", "norm3", "ffn"): + if hasattr(old_block, attr) and hasattr(new_block, attr): + getattr(new_block, attr).load_state_dict(getattr(old_block, attr).state_dict()) + if hasattr(old_block, "modulation") and hasattr(new_block, "modulation"): + with torch.no_grad(): + new_block.modulation.copy_(old_block.modulation.to(dtype=block_dtype)) + # # self_attn: SSE-GLA structure differs from Wan SelfAttention, + # # so remap Wan's {q,k,v,o} -> {q,k,v,o}_proj and load (norm_q/ + # # norm_k have no SSE-GLA counterpart -> dropped). + # _sa_sd = {} + # for _k, _v in old_block.self_attn.state_dict().items(): + # if _k.startswith("q."): + # _sa_sd["q_proj." + _k[2:]] = _v + # elif _k.startswith("k."): + # _sa_sd["k_proj." + _k[2:]] = _v + # elif _k.startswith("v."): + # _sa_sd["v_proj." + _k[2:]] = _v + # elif _k.startswith("o."): + # _sa_sd["o_proj." + _k[2:]] = _v + # new_block.self_attn.load_state_dict(_sa_sd, strict=False) + else: + new_block = DiTBlock_w_Action( + has_image_input=has_image_input, + dim=dim, num_heads=num_heads, ffn_dim=ffn_dim, eps=eps, + add_action_attn=_arg('add_action_attn', False), + action_use_temporal_attention=_arg('action_use_temporal_attention', False), + use_cam_pose=_use_cam_pose, + use_block_wise_ssm=attach_block_ssm, + use_videossm_hybrid=attach_videossm, + videossm_kernel_size=int(_arg('videossm_kernel_size', 3) or 3), + videossm_expand=int(_arg('videossm_expand', 2) or 2), + ) + new_block = new_block.to(dtype=block_dtype, device=next(old_block.parameters()).device) + for attr in ("self_attn", "cross_attn", "norm1", "norm2", "norm3", "ffn"): + if hasattr(old_block, attr) and hasattr(new_block, attr): + getattr(new_block, attr).load_state_dict(getattr(old_block, attr).state_dict()) + if hasattr(old_block, "modulation") and hasattr(new_block, "modulation"): + with torch.no_grad(): + new_block.modulation.copy_(old_block.modulation.to(dtype=block_dtype)) + new_blocks.append(new_block) + + dit.blocks = new_blocks + _n_cgla = sum(1 for b in new_blocks if isinstance(b, CGLATransformerBlock)) + if use_cgla_memory: + _mech = str(_arg('cgla_mechanism', 'cgla') or 'cgla') + logger.info(f"[CGLA] Replaced {_n_cgla}/{len(new_blocks)} DiT blocks with " + f"CGLATransformerBlock (every_n={cgla_every_n}, mechanism={_mech}, " + f"head_dim={_cgla_head_dim}, bidirectional); " + f"Wan self_attn q/k/v/o -> SSE-GLA q/k/v/o_proj") + + _mlp_type = "MLP_CamPose" if _use_cam_pose else "MLP_Action" + logger.info(f"[VWM-style] Replaced {len(new_blocks)} DiT blocks with DiTBlock_w_Action ({_mlp_type}, zero-init)") + if use_block_wise_ssm: + logger.info(f"[Block-wise SSM] attached to every {ssm_every_n} DiT block(s)") + if use_videossm_hybrid: + logger.info(f"[VideoSSM hybrid] attached to every {videossm_every_n} DiT block(s)") + + device = next(dit.parameters()).device + _ckpt_path = _arg('ckpt_path', None) or _arg('resume_from_checkpoint', None) + if _ckpt_path is not None: + ckpt = safe_load_file(_ckpt_path) + if use_cgla_memory: + # Remap Wan self_attn.{q,k,v,o} -> SSE-GLA {q,k,v,o}_proj so a Wan + # checkpoint loads the linear-attention projections. No-op on an + # already-CGLA checkpoint (keys are already *_proj). + ckpt = remap_wan_to_cgla(ckpt) + missing, unexpected = dit.load_state_dict(ckpt, strict=False) + dit.to(device=device) + logger.info(f"[VWM-style] Loaded ckpt: {len(ckpt)} keys, missing={len(missing)}, unexpected={len(unexpected)}") + + if use_cgla_memory: + # Block-aware freeze (supports cgla_every_n_blocks > 1, where some + # blocks stay DiTBlock_w_Action = Wan softmax). CGLA blocks: train + # the SSE-GLA self_attn + cam_encoder + cgla_gate + noise_write_gate, + # freeze the Wan backbone (cross_attn/norm/ffn/modulation). Non-CGLA + # blocks: VWM pattern (train action_mlp / self_attn_with_action / + # SSM, freeze Wan softmax self_attn + backbone). + for block in dit.blocks: + if isinstance(block, CGLATransformerBlock): + for name, param in block.named_parameters(): + if ("action_mlp" in name) or ("self_attn_with_action" in name) or ("self_attn" in name) \ + ("cam_encoder" in name) or ("noise_write_gate" in name): + param.requires_grad = True + else: + param.requires_grad = False + else: + for name, param in block.named_parameters(): + if ("action_mlp" in name) or ("self_attn_with_action" in name) or \ + ("block_wise_ssm" in name) or ("videossm_hybrid" in name): + param.requires_grad = True + else: + param.requires_grad = False + elif _arg('action_module_only', False): + if _arg('add_action_attn', False): + for block in dit.blocks: + for name, param in block.named_parameters(): + if ("action_mlp" in name) or ("self_attn_with_action" in name) or ("block_wise_ssm" in name) or ("videossm_hybrid" in name): + param.requires_grad = True + else: + param.requires_grad = False + else: + for block in dit.blocks: + for name, param in block.named_parameters(): + if "action_mlp" in name or "self_attn" in name or "block_wise_ssm" in name or "videossm_hybrid" in name: + param.requires_grad = True + else: + param.requires_grad = False + else: + for block in dit.blocks: + for name, param in block.named_parameters(): + if "action_mlp" in name or "self_attn_with_action" in name or "block_wise_ssm" in name or "videossm_hybrid" in name: + param.requires_grad = True + else: + param.requires_grad = False + _log_dit_freeze_summary(dit) + + _resume_from = _arg('resume_from', None) + if _resume_from: + logger.info(f"Loading full resume checkpoint: {_resume_from}") + ckpt = safe_load_file(_resume_from) + if use_cgla_memory: + # No-op on an already-CGLA ckpt (keys are *_proj); remaps a Wan ckpt. + ckpt = remap_wan_to_cgla(ckpt) + model.pipe.dit.load_state_dict(ckpt, strict=False) + logger.info(f"Checkpoint loaded, resuming from step {resume_step_count}") + + model_logger = ModelLogger( + args.output_path, + remove_prefix_in_ckpt=args.remove_prefix_in_ckpt, + wandb_run_name=args.wandb_run_name, + ckpt_interval=args.ckpt_interval, + resume_step_count=resume_step_count, + save_full_model=_arg('save_full_model', False), + context_drop_prob=float(_arg("context_drop_prob", 0.0)), + enable_video_sampling=_arg("enable_video_sampling", False), + sampling_interval_steps=int(_arg("sampling_interval_steps", 0)), + sampling_two_chunk_memory=_arg("sampling_two_chunk_memory", False), + sampling_action_path=_arg("sampling_action_path", None), + sampling_two_chunk_action_path=_arg("sampling_two_chunk_action_path", None), + sampling_negative_prompt=_arg("sampling_negative_prompt", ""), + sampling_height=int(_arg("sampling_height", 352)), + sampling_width=int(_arg("sampling_width", 640)), + sampling_num_frames=int(_arg("sampling_num_frames", 81)), + sampling_num_inference_steps=int(_arg("sampling_num_inference_steps", 50)), + context_memory_frames=int(_arg("context_memory_frames", 1)), + context_source=_arg("context_source", "replay"), + context_per_frame_vae=_arg("context_per_frame_vae", False), + # Monitor samples with the SAME noise-schedule shift training uses + # (--timestep_shift), so the in-training video reflects the trained + # schedule (15 for two-chunk rows, 5 for legacy ctx rows). + sampling_sigma_shift=float(_arg("timestep_shift", 1.0) or 1.0), + ) + + optimizer = torch.optim.AdamW(model.trainable_modules(), lr=args.learning_rate) + scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer) + + # Setup FOV retriever for context-based memory training (also for ModelLogger sampling) + enable_fov_retrieval = _arg('enable_fov_retrieval', False) + fov_retriever = None + dataset_base_path = _arg('dataset_base_path', None) + if enable_fov_retrieval: + fov_retriever = setup_fov_retriever_for_training( + dataset_base_path=dataset_base_path, + enable_fov_retrieval=True + ) + + launch_training_task( + dataset, model, model_logger, optimizer, scheduler, + num_epochs=args.num_epochs, + gradient_accumulation_steps=args.gradient_accumulation_steps, + per_device_train_batch_size=int(_arg("per_device_train_batch_size", 1)), + spike_threshold=_arg('spike_threshold', 5.0), + resume_step_count=resume_step_count, + enable_fov_retrieval=enable_fov_retrieval, + retrieval_method=_arg('retrieval_method', 'fov'), + latent_retrieval_dir=_arg('latent_retrieval_dir', None), + dataset_base_path=_arg('dataset_base_path', None), + fov_retriever=fov_retriever, + context_memory_frames=_arg('context_memory_frames', 8), + prev_chunk_frames=int(_arg('prev_chunk_frames', 81)), + fov_top_k=_arg('fov_top_k', 4), # Number of overlap frames (4), GT frame 0 added automatically + use_rt_relative=_arg('use_rt_relative', False), # Experiment 1_4_2: RT relative conversion + strict_overlap_context=_arg('strict_overlap_context', False), + dataset_repeat=_arg('dataset_repeat', 1), # Pass dataset_repeat for step calculation + use_camera_encoder=_arg('use_camera_encoder', False), # exp1_4_3: DDP find_unused_parameters + num_workers=_arg('num_workers', 0), + context_source=_arg('context_source', 'fov'), + max_train_steps=int(_arg('max_train_steps', 0)), + progress_total_steps=int(_arg('progress_total_steps', 0)), + log_interval=max(1, int(_arg('log_interval', 20) or 20)), + max_grad_norm=float(_arg('max_grad_norm', 1.0)), + ) + + model_logger.finish() \ No newline at end of file diff --git a/code/src/model_training/training_modules/__init__.py b/code/src/model_training/training_modules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3ace1b53769225070083b99e63e8edfe8d434d51 --- /dev/null +++ b/code/src/model_training/training_modules/__init__.py @@ -0,0 +1,13 @@ +from .action_blocks import DiTBlock_w_Action, MLP_Action, MLP_CamPose +from .model_logger import ModelLogger +from .training_loop import launch_training_task +from .wan_training_module import WanTrainingModule + +__all__ = [ + "DiTBlock_w_Action", + "MLP_Action", + "MLP_CamPose", + "ModelLogger", + "WanTrainingModule", + "launch_training_task", +] diff --git a/code/src/model_training/training_modules/action_blocks.py b/code/src/model_training/training_modules/action_blocks.py new file mode 100644 index 0000000000000000000000000000000000000000..3553e402937d8f5061a0a774c7d591a7810922c1 --- /dev/null +++ b/code/src/model_training/training_modules/action_blocks.py @@ -0,0 +1,127 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +from diffsynth.models.memory.block_wise_ssm import BlockWiseStateSpaceMemory +from diffsynth.models.memory.videossm_hybrid import HybridStateSpaceMemory +from diffsynth.models.wan_video_dit import SelfAttention, CrossAttention, GateModule, modulate + + +class MLP_Action(nn.Module): + def __init__(self, out_dim, sliding_window_size=3, r=4): + super().__init__() + self.proj_action = nn.Linear(r * sliding_window_size * 10, out_dim) + nn.init.zeros_(self.proj_action.weight) + nn.init.zeros_(self.proj_action.bias) + self.sliding_window_size = sliding_window_size + self.r = r + + def forward(self, x): + bs, nr, act_dim = x.shape + r = self.r + n = nr // r + actions = x.reshape(bs, n, r, act_dim) + actions = F.pad(actions, (0, 0, 0, 0, self.sliding_window_size - 1, 1), mode="replicate") + action_windows = [] + for i in range(self.sliding_window_size): + action_windows.append(actions[:, i:i + n + 1]) + actions = torch.cat(action_windows, dim=2) + actions = actions.reshape(bs, n + 1, -1) + actions = self.proj_action(actions) + return actions + + +class MLP_CamPose(nn.Module): + def __init__(self, out_dim, pose_dim=12): + super().__init__() + self.proj = nn.Linear(pose_dim, out_dim) + nn.init.zeros_(self.proj.weight) + nn.init.zeros_(self.proj.bias) + + def forward(self, x): + return self.proj(x) + + +class DiTBlock_w_Action(nn.Module): + def __init__(self, has_image_input: bool, dim: int, num_heads: int, ffn_dim: int, + eps: float = 1e-6, add_action_attn=False, + action_use_temporal_attention: bool = True, use_cam_pose: bool = False, + use_block_wise_ssm: bool = False, use_videossm_hybrid: bool = False, + videossm_kernel_size: int = 3, videossm_expand: int = 2): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.ffn_dim = ffn_dim + + if add_action_attn: + self.self_attn_with_action = SelfAttention(dim, num_heads, eps) + nn.init.zeros_(self.self_attn_with_action.o.weight) + nn.init.zeros_(self.self_attn_with_action.o.bias) + if use_cam_pose: + self.action_mlp = MLP_CamPose(dim) + else: + self.action_mlp = MLP_Action(dim) + + self.self_attn = SelfAttention(dim, num_heads, eps) + self.cross_attn = CrossAttention(dim, num_heads, eps, has_image_input=has_image_input) + self.norm1 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) + self.norm2 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) + self.norm3 = nn.LayerNorm(dim, eps=eps) + self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'), nn.Linear(ffn_dim, dim)) + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + self.gate = GateModule() + self.action_use_temporal_attention = action_use_temporal_attention + self.use_block_wise_ssm = bool(use_block_wise_ssm) + self.use_videossm_hybrid = bool(use_videossm_hybrid) + if use_block_wise_ssm: + self.block_wise_ssm = BlockWiseStateSpaceMemory(dim) + if use_videossm_hybrid: + self.videossm_hybrid = HybridStateSpaceMemory( + dim, kernel_size=videossm_kernel_size, expand=videossm_expand + ) + + def forward(self, x, context, t_mod, freqs, actions=None): + has_seq = len(t_mod.shape) == 4 + chunk_dim = 2 if has_seq else 1 + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(6, dim=chunk_dim) + if has_seq: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + shift_msa.squeeze(2), scale_msa.squeeze(2), gate_msa.squeeze(2), + shift_mlp.squeeze(2), scale_mlp.squeeze(2), gate_mlp.squeeze(2), + ) + + num_frames = None + if actions is not None: + original_x = x + actions = self.action_mlp(actions.to(x.dtype)).to(x.dtype) + bs, num_frames, dim = actions.shape + actions = actions.reshape(bs, num_frames, 1, dim) + x = x.reshape(bs, num_frames, -1, dim) + x = x + actions + if hasattr(self, "self_attn_with_action"): + if not self.action_use_temporal_attention: + x = x.reshape(bs, -1, dim) + x = original_x + self.self_attn_with_action(x, freqs) + else: + from einops import rearrange + x = rearrange(x, "b f p d -> (b p) f d") + attn_out = self.self_attn_with_action(x) + attn_out = rearrange(attn_out, "(b p) f d -> b f p d", b=bs) + x = original_x + attn_out.reshape(bs, -1, dim) + else: + x = x.reshape(bs, -1, dim) + + input_x = modulate(self.norm1(x), shift_msa, scale_msa) + x = self.gate(x, gate_msa, self.self_attn(input_x, freqs)) + if num_frames is not None: + if hasattr(self, "block_wise_ssm"): + x = self.block_wise_ssm(x, f=num_frames) + if hasattr(self, "videossm_hybrid"): + spatial = x.shape[1] // int(num_frames) if int(num_frames) > 0 else 0 + x = self.videossm_hybrid(x, f=num_frames, h=1, w=spatial) + x = x + self.cross_attn(self.norm3(x), context) + input_x = modulate(self.norm2(x), shift_mlp, scale_mlp) + x = self.gate(x, gate_mlp, self.ffn(input_x)) + return x + diff --git a/code/src/model_training/training_modules/model_logger.py b/code/src/model_training/training_modules/model_logger.py new file mode 100644 index 0000000000000000000000000000000000000000..ff269be95a71e4b28d1b82cea5cafbcb442bf584 --- /dev/null +++ b/code/src/model_training/training_modules/model_logger.py @@ -0,0 +1,211 @@ +import json +import math +import os +from typing import Optional + +import wandb + +from src.model_training.transformers_compat import patch_transformers_hybrid_cache + +patch_transformers_hybrid_cache() + +from diffsynth.trainers.utils import ModelLogger as BaseModelLogger + + +class ModelLogger(BaseModelLogger): + """Compatibility wrapper for legacy training scripts.""" + + def __init__( + self, + output_path, + remove_prefix_in_ckpt=None, + state_dict_converter=lambda x: x, + wandb_run_name=None, + ckpt_interval=None, + resume_step_count=0, + save_full_model=False, + context_drop_prob: float = 0.0, + enable_video_sampling=False, + sampling_interval_steps: int = 0, + sampling_two_chunk_memory: bool = False, + sampling_action_path: Optional[str] = None, + sampling_two_chunk_action_path: Optional[str] = None, + sampling_negative_prompt: str = "oversaturated colors, overexposed, static, blurry details", + sampling_height: int = 352, + sampling_width: int = 640, + sampling_num_frames: int = 81, + sampling_num_inference_steps: int = 50, + context_memory_frames: int = 1, + context_source: str = "replay", + context_per_frame_vae: bool = False, + sampling_sigma_shift: float = 5.0, + ): + super().__init__(output_path, remove_prefix_in_ckpt=remove_prefix_in_ckpt, state_dict_converter=state_dict_converter) + self.wandb_run_name = wandb_run_name + self.ckpt_interval = int(ckpt_interval) if ckpt_interval else None + self.step_count = int(resume_step_count) + self.save_full_model = bool(save_full_model) + self.total_steps = None + self.context_drop_prob = float(context_drop_prob) + self.enable_video_sampling = bool(enable_video_sampling) + self.sampling_interval_steps = int(sampling_interval_steps) + self.sampling_two_chunk_memory = bool(sampling_two_chunk_memory) + self.sampling_action_path = sampling_action_path + self.sampling_two_chunk_action_path = sampling_two_chunk_action_path + self.sampling_negative_prompt = sampling_negative_prompt + self.sampling_height = int(sampling_height) + self.sampling_width = int(sampling_width) + self.sampling_num_frames = int(sampling_num_frames) + self.sampling_num_inference_steps = int(sampling_num_inference_steps) + self.context_memory_frames = int(context_memory_frames) + self.context_source = context_source.strip().lower() + self.context_per_frame_vae = bool(context_per_frame_vae) + # Noise-schedule shift used by the periodic sampling monitor. Must match + # the TRAINING shift (set in train.py from --timestep_shift), else the + # monitor samples from a noise schedule the model never trained on -> + # structure roughly OK but colors blown out / distorted. + self.sampling_sigma_shift = float(sampling_sigma_shift) + self.wandb_logger = None + if self.wandb_run_name: + self.wandb_logger = wandb.init(project="wan-cam", name=self.wandb_run_name, reinit=True) + + def _save_step_or_epoch_ckpt(self, accelerator, model, path: str): + state_dict = None + unwrapped = accelerator.unwrap_model(model) + if self.save_full_model: + # Save full DiT (including action/camera/memory modules), not whole pipeline. + state_dict = accelerator.get_state_dict(unwrapped.pipe.dit) + for module_name in ( + "spatial_memory_module", + "spatial_memory_readout_module", + "geometry_spatial_memory_module", + "geometry_spatial_memory_readout_module", + ): + module = getattr(unwrapped, module_name, None) + if module is not None: + state_dict.update( + { + f"{module_name}.{name}": param + for name, param in accelerator.get_state_dict(module).items() + } + ) + if state_dict is None: + full_state = accelerator.get_state_dict(model) + state_dict = unwrapped.export_trainable_state_dict(full_state, remove_prefix=self.remove_prefix_in_ckpt) + state_dict = self.state_dict_converter(state_dict) + os.makedirs(self.output_path, exist_ok=True) + accelerator.save(state_dict, path, safe_serialization=True) + + def _maybe_sample_paper_process(self, accelerator=None, model=None, current_batch=None): + if not ( + self.enable_video_sampling + and self.sampling_two_chunk_memory + and self.sampling_interval_steps > 0 + and self.step_count % self.sampling_interval_steps == 0 + and accelerator is not None + and model is not None + and current_batch is not None + ): + return + from diffsynth import save_video + from src.model_training.multichunk_sample_utils import ( + run_two_chunk_memory_monitor, + sync_pipe_memory_from_training_module, + ) + + sample = current_batch[0] if isinstance(current_batch, list) else current_batch + first_frame = sample["video"][0] + unwrapped = accelerator.unwrap_model(model) + pipe = unwrapped.pipe + sync_pipe_memory_from_training_module(pipe, unwrapped) + action0 = self.sampling_two_chunk_action_path or self.sampling_action_path + action1 = self.sampling_action_path + # Sampling runs the inference pipe.__call__, which (a) overwrites the + # training scheduler with an inference schedule and (b) spikes GPU + # memory on the main rank. Wrap it so a sampling failure can never abort + # training, and always restore training scheduler state + free cache. + try: + frames0, frames1, meta = run_two_chunk_memory_monitor( + pipe, + prompt=sample.get("prompt") or sample.get("description") or "A scene.", + negative_prompt=self.sampling_negative_prompt, + action_path=self.sampling_action_path, + chunk0_action_path=action0, + chunk1_action_path=action1, + first_frame_pil=first_frame, + context_memory_frames=self.context_memory_frames, + chunk_frames=self.sampling_num_frames, + h=self.sampling_height, + w=self.sampling_width, + seed=42 + self.step_count + accelerator.process_index, + sigma_shift=self.sampling_sigma_shift, + num_inference_steps=self.sampling_num_inference_steps, + cfg_scale=5.0, + inference_noise_level=0.0, + omit_context_actions=False, + context_source=self.context_source, + context_position=os.environ.get("CONTEXT_POSITION", "suffix"), + context_per_frame_vae=self.context_per_frame_vae, + device=pipe.device, + log_prefix=f"[paper-sampling][step={self.step_count}]", + ) + out_dir = os.path.join(self.output_path, "paper_process_sampling") + os.makedirs(out_dir, exist_ok=True) + tag = f"step_{self.step_count:07d}_rank{accelerator.process_index}" + save_video(list(frames0) + list(frames1), os.path.join(out_dir, f"{tag}_pred.mp4"), fps=15, quality=5) + with open(os.path.join(out_dir, f"{tag}_meta.json"), "w", encoding="utf-8") as f: + json.dump(meta, f, ensure_ascii=False, indent=2) + except Exception as e: + # Sampling is a monitoring side-effect; never let it kill training. + import traceback + print(f"[paper-sampling][step={self.step_count}] FAILED (training continues): {e}") + traceback.print_exc() + finally: + # Restore training scheduler + free the sampling step's GPU cache on + # the main rank so the next training step (target VAE encode) is clean. + _train_module = getattr(unwrapped, "restore_after_sampling", None) + if callable(_train_module): + _train_module() + else: + import gc + gc.collect() + if __import__("torch").cuda.is_available(): + __import__("torch").cuda.empty_cache() + + def on_step_end(self, loss, accelerator=None, model=None, current_batch=None, + aux_loss=None): + self.step_count += 1 + if self.wandb_logger is not None: + if accelerator is None or accelerator.is_main_process: + loss_v = float(loss.detach().float().item()) + metrics = {"train/loss": loss_v, "step": self.step_count} + # CGLA SSE load-balancing aux loss (summed per-block in the + # training loop; passed in so dummy/skip steps don't log stale + # values). Logged only when CGLA is active (cgla_aux_loss_weight>0). + if aux_loss is not None and not math.isnan(aux_loss): + metrics["train/aux"] = float(aux_loss) + self.wandb_logger.log(metrics) + if accelerator is not None and accelerator.is_main_process: + self._maybe_sample_paper_process(accelerator, model, current_batch) + if accelerator is not None and self.enable_video_sampling and self.sampling_two_chunk_memory and self.sampling_interval_steps > 0: + accelerator.wait_for_everyone() + if ( + self.ckpt_interval + and accelerator is not None + and model is not None + and (self.step_count % self.ckpt_interval == 0) + ): + accelerator.wait_for_everyone() + if accelerator.is_main_process: + path = os.path.join(self.output_path, f"Step-{self.step_count}.safetensors") + self._save_step_or_epoch_ckpt(accelerator, model, path) + + def on_epoch_end(self, accelerator, model, epoch_id): + accelerator.wait_for_everyone() + if accelerator.is_main_process: + path = os.path.join(self.output_path, f"epoch-{epoch_id}.safetensors") + self._save_step_or_epoch_ckpt(accelerator, model, path) + + def finish(self): + if self.wandb_logger is not None: + wandb.finish() diff --git a/code/src/model_training/training_modules/training_loop.py b/code/src/model_training/training_modules/training_loop.py new file mode 100644 index 0000000000000000000000000000000000000000..f51a9e63ca1676a5f588ca7d680e17a628955099 --- /dev/null +++ b/code/src/model_training/training_modules/training_loop.py @@ -0,0 +1,443 @@ +import logging +import math +import os +from typing import Optional + +import torch +import torch.distributed as dist +from accelerate import Accelerator +from tqdm import tqdm + +from src.model_training.transformers_compat import patch_transformers_hybrid_cache + +patch_transformers_hybrid_cache() + +from diffsynth.trainers.utils import DiffusionTrainingModule +from src.model_training.fov_retrieval import FOVMemoryRetriever +from src.model_training.fov_retrieval import retrieve_context_frames_advanced, retrieve_fov_context_frames +from src.model_training.training_modules.model_logger import ModelLogger + +logger = logging.getLogger(__name__) + + +def launch_training_task( + dataset: torch.utils.data.Dataset, + model: DiffusionTrainingModule, + model_logger: ModelLogger, + optimizer: torch.optim.Optimizer, + scheduler: torch.optim.lr_scheduler.LRScheduler, + num_epochs: int = 1, + gradient_accumulation_steps: int = 1, + per_device_train_batch_size: int = 1, + seed: int = 42, + spike_threshold: float = 5.0, + resume_step_count: int = 0, + enable_fov_retrieval: bool = False, + retrieval_method: str = "fov", # fov | latent_sim + latent_retrieval_dir: Optional[str] = None, + dataset_base_path: str = None, + fov_retriever: Optional[FOVMemoryRetriever] = None, + context_memory_frames: int = 5, + prev_chunk_frames: int = 81, + fov_top_k: int = 4, # Number of overlap frames to retrieve. GT frame 0 will be added automatically. + use_rt_relative: bool = False, # Experiment 1_4_2: Use RT relative conversion (aligned with Context-as-Memory) + strict_overlap_context: bool = False, + dataset_repeat: int = 1, # Add dataset_repeat parameter for step calculation + use_camera_encoder: bool = False, # exp1_4_3: use CameraEncoder (action_mlp unused -> need find_unused_parameters) + num_workers: int = 0, # DataLoader workers: 0=main process, >0=parallel preload (recommend 4 for video) + context_source: str = "fov", + max_train_steps: int = 0, + progress_total_steps: int = 0, + log_interval: int = 20, + max_grad_norm: float = 1.0, +): + prev_chunk_frames = int(prev_chunk_frames) + # VideoDataset can return None when file loading fails; keep distributed batches aligned. + def collate_fn(batch): + valid_batch = [item for item in batch if item is not None] + return valid_batch or None + + num_workers = max(0, int(num_workers)) + dataloader = torch.utils.data.DataLoader( + dataset, + batch_size=per_device_train_batch_size, + shuffle=True, + collate_fn=collate_fn, + num_workers=num_workers, + drop_last=True, + persistent_workers=(num_workers > 0), + pin_memory=(num_workers > 0 and torch.cuda.is_available()), + ) + if num_workers > 0: + logger.info(f"[DataLoader] num_workers={num_workers}, persistent_workers=True, pin_memory={torch.cuda.is_available()} (data preload parallel to GPU)") + + timeout_seconds = int(os.environ.get('TORCH_DISTRIBUTED_DEFAULT_TIMEOUT', 2400)) + os.environ['TORCH_DISTRIBUTED_DEFAULT_TIMEOUT'] = str(timeout_seconds) + logger.info(f"[Timeout Config] Setting TORCH_DISTRIBUTED_DEFAULT_TIMEOUT={timeout_seconds} seconds ({timeout_seconds/60:.1f} minutes)") + + # Conditional context paths can leave parameters unused on some iterations. + need_find_unused = bool(use_camera_encoder) or model_logger.context_drop_prob > 0.0 + if need_find_unused: + from accelerate import DistributedDataParallelKwargs + ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) + accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps, kwargs_handlers=[ddp_kwargs]) + logger.info("[DDP] find_unused_parameters=True (conditional modules / context_drop_prob enabled)") + else: + accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps) + model, optimizer, dataloader, scheduler = accelerator.prepare(model, optimizer, dataloader, scheduler) + + if model_logger.enable_video_sampling and model_logger.total_steps is not None: + dataset_size = len(dataset) + num_processes = accelerator.num_processes + effective_dataset_size = dataset_size * dataset_repeat + total_steps_per_gpu = (effective_dataset_size * num_epochs) // (gradient_accumulation_steps * num_processes * per_device_train_batch_size) + total_steps_global = total_steps_per_gpu * num_processes + model_logger.total_steps = total_steps_global + + if accelerator.is_main_process: + logger.info("="*80) + logger.info("[Step Calculation] Corrected total_steps after accelerator.init") + logger.info("="*80) + logger.info(f" Dataset size (unique samples): {dataset_size}") + logger.info(f" Dataset repeat: {dataset_repeat}") + logger.info(f" Effective dataset size: {effective_dataset_size} (unique * repeat)") + logger.info(f" Number of epochs: {num_epochs}") + logger.info(f" Number of GPUs: {num_processes}") + logger.info(f" Gradient accumulation steps: {gradient_accumulation_steps}") + logger.info(f" Per-device batch size: {per_device_train_batch_size}") + logger.info(f" Total samples to process: {effective_dataset_size * num_epochs}") + logger.info(f" Steps per GPU: ~{total_steps_per_gpu}") + logger.info(f" Total steps (global): {total_steps_global}") + logger.info("") + logger.info(f" ✓ Each GPU will process ~{total_steps_per_gpu} steps") + logger.info(f" ✓ This ensures traversal of all {effective_dataset_size} samples") + logger.info(f" ({dataset_size} unique samples × {dataset_repeat} repeats)") + logger.info(f" ✓ Over {num_epochs} epoch(s)") + logger.info("="*80) + + step = resume_step_count + traj_loss = 0.0 + if resume_step_count > 0: + adaptation_steps = max(200, resume_step_count // 100) + spike_detection_start_step = resume_step_count + adaptation_steps + logger.info(f"Resuming from step {resume_step_count}, spike detection will start at step {spike_detection_start_step} (after {adaptation_steps} adaptation steps)") + else: + spike_detection_start_step = 100 + + for epoch_id in range(num_epochs): + epoch_seed = seed + epoch_id + torch.manual_seed(epoch_seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed(epoch_seed) + torch.cuda.manual_seed_all(epoch_seed) + + if resume_step_count > 0 and epoch_id == 0: + estimated_skip = resume_step_count // gradient_accumulation_steps + if estimated_skip > 0: + logger.info(f"Skipping {estimated_skip} data samples to resume from step {resume_step_count}...") + dataloader_iter = iter(dataloader) + for _ in tqdm(range(estimated_skip), desc="Skipping data", unit="samples", leave=False): + try: + next(dataloader_iter) + except StopIteration: + break + dataloader = dataloader_iter + logger.info(f"Successfully skipped {estimated_skip} data samples, resuming training...") + + # Track consecutive None data to detect if we're stuck in a loop + consecutive_none_count = 0 + max_consecutive_none = 100 # If we get 100 consecutive None values, something is wrong + + progress_total = int(progress_total_steps) + if progress_total <= 0: + progress_total = len(dataloader) + progress_bar = tqdm( + dataloader, + total=progress_total, + initial=resume_step_count if progress_total_steps else 0, + desc="Training steps", + unit="step", + ) + for data_idx, data in enumerate(progress_bar): + # Handle None data (can happen if all files in batch fail to load) + if data is None: + consecutive_none_count += 1 + if consecutive_none_count >= max_consecutive_none: + logger.error(f"Received {max_consecutive_none} consecutive None data samples. This suggests a serious dataset issue. Stopping training.") + raise ValueError(f"Too many consecutive None data samples ({max_consecutive_none}). Check dataset files.") + + # Log warning but continue (will skip this step) + if consecutive_none_count <= 10 or consecutive_none_count % 10 == 0: + logger.warning(f"Received None data at index {data_idx} (consecutive: {consecutive_none_count}). This may indicate missing or corrupted files. Skipping...") + + # Still increment step to keep step_count synchronized + step += 1 + dummy_loss = torch.tensor(0.0, device=accelerator.device, requires_grad=False) + model_logger.on_step_end(dummy_loss, accelerator, model, current_batch=samples) + continue + + # Reset consecutive None counter when we get valid data + consecutive_none_count = 0 + + # Normalize to list of samples for batch processing (per_device_train_batch_size > 1) + samples = data if isinstance(data, list) else [data] + + # Simplified context-based retrieval OR replay/prev_chunk_tail (aligned with multichunk eval) + context_retrieval_success = True # Set False if any sample fails (for strict mode) + _umodel = accelerator.unwrap_model(model) + _cm_frames = int(_umodel.context_memory_frames) + _cs = context_source.strip().lower() + if _cs not in ("fov", "replay", "prev_chunk_tail"): + _cs = "fov" + + if _cs == "replay" and dataset_base_path: + from src.model_training.multichunk_sample_utils import ( + replay_context_actions_from_segment_actions, + replay_context_global_indices, + synthetic_replay_context_from_segment, + ) + for d in samples: + vf = d.get("video") or [] + n_seg = min(int(prev_chunk_frames), len(vf)) if vf else 0 + ctx_pil = synthetic_replay_context_from_segment(vf, n_seg, _cm_frames) if n_seg > 0 else None + if not ctx_pil: + context_retrieval_success = False + break + d["context_frames"] = ctx_pil + d["context_source"] = "replay_synthetic" + acts = d.get("actions") + if isinstance(acts, list) and len(acts) >= n_seg: + ra = replay_context_actions_from_segment_actions(acts[:n_seg], n_seg, _cm_frames) + if ra is not None: + d["context_actions"] = ra + sf = int(d.get("start_frame", 0) or 0) + idxs = replay_context_global_indices(n_seg, _cm_frames) + d["context_frame_indices"] = [sf + int(i) for i in idxs] + + elif _cs == "prev_chunk_tail" and dataset_base_path: + from src.model_training.multichunk_sample_utils import load_prev_chunk_tail_from_disk, load_prev_chunk_tail_rt_actions + _ctx_pos = os.environ.get("CONTEXT_POSITION", "suffix").strip().lower() + _nearest_first = (_ctx_pos == "suffix") + for d in samples: + sf = int(d.get("start_frame", 0) or 0) + vn = d.get("video_name", "") + pil_list, idxs = load_prev_chunk_tail_from_disk( + dataset_base_path, str(vn), sf, _cm_frames, nearest_first=_nearest_first + ) + if not pil_list: + context_retrieval_success = False + break + d["context_frames"] = pil_list + d["context_frame_indices"] = list(idxs) if idxs else [] + d["context_source"] = "prev_chunk_tail" + ra, _ = load_prev_chunk_tail_rt_actions( + dataset_base_path, + str(vn), + sf, + _cm_frames, + use_rt_relative=use_rt_relative, + nearest_first=_nearest_first, + ) + if ra: + d["context_actions"] = ra + + elif enable_fov_retrieval and dataset_base_path: + for d in samples: + if retrieval_method == "latent_sim": + ( + context_frames, + context_actions, + context_indices, + ref_frame_idx, + video_name, + source, + ) = retrieve_context_frames_advanced( + data=d, + dataset_base_path=dataset_base_path, + top_k=fov_top_k, + drop_overlap_probability=0.1, + use_rt_relative=use_rt_relative, + retrieval_method="latent_sim", + latent_retrieval_dir=latent_retrieval_dir, + strict_overlap_labels=strict_overlap_context, + ) + else: + ( + context_frames, + context_actions, + context_indices, + ref_frame_idx, + video_name, + source, + ) = retrieve_fov_context_frames( + data=d, + dataset_base_path=dataset_base_path, + fov_retriever=fov_retriever, # unused in simplified retrieval, kept for compat + top_k=fov_top_k, # fov_top_k is number of overlap frames (4), GT frame 0 will be added automatically + use_precomputed_overlaps=True, + strict_overlap_labels=strict_overlap_context, + allow_realtime_fallback=(not strict_overlap_context), + allow_segment_fallback=(not strict_overlap_context), + ) + + if context_frames and len(context_frames) > 0: + # Use retrieved frames as context + d["context_frames"] = context_frames + if context_actions: + d["context_actions"] = context_actions + # Store retrieval metadata for visualization/debugging + d["context_frame_indices"] = context_indices + d["context_ref_frame_idx"] = ref_frame_idx + d["context_video_name"] = video_name + d["context_source"] = source + else: + context_retrieval_success = False + break + + # Strict mode: if we require context but retrieval failed, skip this step + _need_ctx_strict = ( + strict_overlap_context + and (not context_retrieval_success) + and ( + enable_fov_retrieval + or context_source.strip().lower() in ("replay", "prev_chunk_tail") + ) + ) + if _need_ctx_strict: + if step % 50 == 0 and accelerator.is_main_process: + logger.warning(f"[CONTEXT][STRICT] No context at step={step}, skipping this training sample.") + step += 1 + dummy_loss = torch.tensor(0.0, device=accelerator.device, requires_grad=False) + model_logger.on_step_end(dummy_loss, accelerator, model, current_batch=samples) + continue + + # Diagnostic: cache the raw dataset video + a VAE encode→decode roundtrip + # every CACHE_INPUT_VIDEO_EVERY steps (env, default 0=off) to check whether + # the videos the model trains on have color issues (inversion / R-B swap). + # Main-process only; no_grad side-effect — never affects the forward/grads. + _cache_every = int(os.environ.get("CACHE_INPUT_VIDEO_EVERY", "0") or "0") + if _cache_every > 0 and accelerator.is_main_process and (step % _cache_every == 0): + try: + _dbg_dir = os.path.join(model_logger.output_path, "input_video_debug") + _umodel.dump_input_video_debug(samples, step, _dbg_dir) + except Exception as _e: # noqa: BLE001 - diagnostic only + print(f"[input-video-debug] step={step} skipped: {_e}", flush=True) + + with accelerator.accumulate(model): + optimizer.zero_grad() + # One forward over full batch: data is list of B dicts when per_device_train_batch_size > 1 + # Main loss on current batch + loss = model(data) + + step += 1 + # nan/inf guard: one non-finite loss otherwise poisons traj_loss (EMA) + # and every subsequent step (nan > threshold is False, so the spike-skip + # below never catches it). Skip synchronously across ranks to avoid DDP hang. + _finite = torch.tensor( + 1.0 if bool(torch.isfinite(loss).all().item()) else 0.0, + device=accelerator.device, dtype=torch.float32) + if accelerator.num_processes > 1: + dist.all_reduce(_finite, op=dist.ReduceOp.MIN) + if _finite.item() < 0.5: + if accelerator.is_main_process: + logger.warning(f"Non-finite loss at step {step}, skipping step (all ranks).") + model_logger.on_step_end(torch.tensor(0.0, device=accelerator.device), accelerator, model, current_batch=samples) + del loss + torch.cuda.empty_cache() + continue + if traj_loss == 0.0: + traj_loss = loss.item() + else: + alpha = 0.01 + traj_loss = (1 - alpha) * traj_loss + alpha * loss.item() + + if step >= spike_detection_start_step and traj_loss > 0: + relative_loss = loss.item() / traj_loss + if resume_step_count > 0 and step < resume_step_count + 500: + effective_threshold = spike_threshold * 1.5 + else: + effective_threshold = spike_threshold + + should_skip = relative_loss > effective_threshold + # Keep the skip decision identical across ranks to avoid DDP hangs. + skip_t = torch.tensor(1.0 if should_skip else 0.0, device=accelerator.device, dtype=torch.float32) + if accelerator.num_processes > 1: + dist.all_reduce(skip_t, op=dist.ReduceOp.MAX) + skip_global = skip_t.item() > 0.5 + + if skip_global: + if accelerator.is_main_process: + logger.warning(f"Spike detected at step {step} (loss={loss.item():.4f}, traj_loss={traj_loss:.4f}, ratio={relative_loss:.2f}), sync skip across all ranks") + dummy_loss = torch.tensor(0.0, device=accelerator.device, requires_grad=False) + model_logger.on_step_end(dummy_loss, accelerator, model, current_batch=samples) + del loss + torch.cuda.empty_cache() + continue + accelerator.backward(loss) + loss_v = float(loss.detach().float().item()) + # CGLA SSE load-balancing aux loss (only when + # cgla_aux_loss_weight > 0). Each CGLATransformerBlock stashes + # its aux_loss during the forward; sum across blocks (mirrors + # WanVideoPipeline.training_loss). Reported for monitoring — it + # is already added to `loss` inside the pipeline. + aux_v = float("nan") + _aux_w = 0.0 + _umodel = accelerator.unwrap_model(model) + _pipe = getattr(_umodel, "pipe", None) + if _pipe is not None: + _aux_w = float(getattr(_pipe, "cgla_aux_loss_weight", 0.0) or 0.0) + if _aux_w > 0: + _aux_sum = None + for _b in getattr(getattr(_pipe, "dit", None), "blocks", []): + _a = getattr(_b, "_aux_loss", None) + if _a is not None and torch.is_tensor(_a) and torch.isfinite(_a).all(): + _aux_sum = _a if _aux_sum is None else _aux_sum + _a + if _aux_sum is not None: + aux_v = float(_aux_sum.detach().float().item()) + # grad_norm / clip only on the accumulation boundary (sync step): + # avoid a wasteful all-reduce every micro-step, and measure the + # full accumulated gradient rather than a partial one. + grad_norm = float("nan") + cur_lr = float(optimizer.param_groups[0].get("lr", 0.0)) + if accelerator.sync_gradients: + if max_grad_norm and max_grad_norm > 0: + grad_norm_t = accelerator.clip_grad_norm_(model.parameters(), max_grad_norm) + else: + grad_norm_t = accelerator.clip_grad_norm_(model.parameters(), float("inf")) + grad_norm = float(grad_norm_t.detach().float().item()) + cur_lr = float(optimizer.param_groups[0].get("lr", 0.0)) + optimizer.step() + # Live progress bar: loss every step; grad_norm/lr on the sync step; + # aux when present. + _post = {"loss": f"{loss_v:.4f}", "traj": f"{traj_loss:.4f}"} + if not math.isnan(grad_norm): + _post["gnorm"] = f"{grad_norm:.3f}" + _post["lr"] = f"{cur_lr:.2e}" + if not math.isnan(aux_v): + _post["aux"] = f"{aux_v:.4f}" + progress_bar.set_postfix(_post) + # Periodic structured line (main process, sync step only) so + # loss/grad_norm/aux land in the tee'd log file, not just the bar. + if (accelerator.sync_gradients and accelerator.is_main_process + and (step % log_interval == 0)): + _aux_str = (f" aux={aux_v:.4f}(w={_aux_w:.3f})" + if not math.isnan(aux_v) else "") + logger.info( + f"[train] step={step} loss={loss_v:.4f} " + f"traj={traj_loss:.4f} grad_norm={grad_norm:.4f} " + f"lr={cur_lr:.2e}{_aux_str}" + ) + model_logger.on_step_end(loss, accelerator, model, current_batch=samples, + aux_loss=aux_v) + scheduler.step() + + if max_train_steps and step >= max_train_steps: + if progress_total_steps: + progress_bar.n = min(step, progress_bar.total) if progress_bar.total is not None else step + progress_bar.refresh() + if accelerator.is_main_process: + logger.info(f"[TRAIN] Reached max_train_steps={max_train_steps}; stopping without epoch checkpoint.") + accelerator.wait_for_everyone() + return + + model_logger.on_epoch_end(accelerator, model, epoch_id) diff --git a/code/src/model_training/training_modules/training_loop.py.nanguard-bak b/code/src/model_training/training_modules/training_loop.py.nanguard-bak new file mode 100644 index 0000000000000000000000000000000000000000..e0cccb339a576c7017444e075f375b0e7c40ff5c --- /dev/null +++ b/code/src/model_training/training_modules/training_loop.py.nanguard-bak @@ -0,0 +1,428 @@ +import logging +import math +import os +from typing import Optional + +import torch +import torch.distributed as dist +from accelerate import Accelerator +from tqdm import tqdm + +from src.model_training.transformers_compat import patch_transformers_hybrid_cache + +patch_transformers_hybrid_cache() + +from diffsynth.trainers.utils import DiffusionTrainingModule +from src.model_training.fov_retrieval import FOVMemoryRetriever +from src.model_training.fov_retrieval import retrieve_context_frames_advanced, retrieve_fov_context_frames +from src.model_training.training_modules.model_logger import ModelLogger + +logger = logging.getLogger(__name__) + + +def launch_training_task( + dataset: torch.utils.data.Dataset, + model: DiffusionTrainingModule, + model_logger: ModelLogger, + optimizer: torch.optim.Optimizer, + scheduler: torch.optim.lr_scheduler.LRScheduler, + num_epochs: int = 1, + gradient_accumulation_steps: int = 1, + per_device_train_batch_size: int = 1, + seed: int = 42, + spike_threshold: float = 5.0, + resume_step_count: int = 0, + enable_fov_retrieval: bool = False, + retrieval_method: str = "fov", # fov | latent_sim + latent_retrieval_dir: Optional[str] = None, + dataset_base_path: str = None, + fov_retriever: Optional[FOVMemoryRetriever] = None, + context_memory_frames: int = 5, + prev_chunk_frames: int = 81, + fov_top_k: int = 4, # Number of overlap frames to retrieve. GT frame 0 will be added automatically. + use_rt_relative: bool = False, # Experiment 1_4_2: Use RT relative conversion (aligned with Context-as-Memory) + strict_overlap_context: bool = False, + dataset_repeat: int = 1, # Add dataset_repeat parameter for step calculation + use_camera_encoder: bool = False, # exp1_4_3: use CameraEncoder (action_mlp unused -> need find_unused_parameters) + num_workers: int = 0, # DataLoader workers: 0=main process, >0=parallel preload (recommend 4 for video) + context_source: str = "fov", + max_train_steps: int = 0, + progress_total_steps: int = 0, + log_interval: int = 20, + max_grad_norm: float = 1.0, +): + prev_chunk_frames = int(prev_chunk_frames) + # VideoDataset can return None when file loading fails; keep distributed batches aligned. + def collate_fn(batch): + valid_batch = [item for item in batch if item is not None] + return valid_batch or None + + num_workers = max(0, int(num_workers)) + dataloader = torch.utils.data.DataLoader( + dataset, + batch_size=per_device_train_batch_size, + shuffle=True, + collate_fn=collate_fn, + num_workers=num_workers, + drop_last=True, + persistent_workers=(num_workers > 0), + pin_memory=(num_workers > 0 and torch.cuda.is_available()), + ) + if num_workers > 0: + logger.info(f"[DataLoader] num_workers={num_workers}, persistent_workers=True, pin_memory={torch.cuda.is_available()} (data preload parallel to GPU)") + + timeout_seconds = int(os.environ.get('TORCH_DISTRIBUTED_DEFAULT_TIMEOUT', 2400)) + os.environ['TORCH_DISTRIBUTED_DEFAULT_TIMEOUT'] = str(timeout_seconds) + logger.info(f"[Timeout Config] Setting TORCH_DISTRIBUTED_DEFAULT_TIMEOUT={timeout_seconds} seconds ({timeout_seconds/60:.1f} minutes)") + + # Conditional context paths can leave parameters unused on some iterations. + need_find_unused = bool(use_camera_encoder) or model_logger.context_drop_prob > 0.0 + if need_find_unused: + from accelerate import DistributedDataParallelKwargs + ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) + accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps, kwargs_handlers=[ddp_kwargs]) + logger.info("[DDP] find_unused_parameters=True (conditional modules / context_drop_prob enabled)") + else: + accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps) + model, optimizer, dataloader, scheduler = accelerator.prepare(model, optimizer, dataloader, scheduler) + + if model_logger.enable_video_sampling and model_logger.total_steps is not None: + dataset_size = len(dataset) + num_processes = accelerator.num_processes + effective_dataset_size = dataset_size * dataset_repeat + total_steps_per_gpu = (effective_dataset_size * num_epochs) // (gradient_accumulation_steps * num_processes * per_device_train_batch_size) + total_steps_global = total_steps_per_gpu * num_processes + model_logger.total_steps = total_steps_global + + if accelerator.is_main_process: + logger.info("="*80) + logger.info("[Step Calculation] Corrected total_steps after accelerator.init") + logger.info("="*80) + logger.info(f" Dataset size (unique samples): {dataset_size}") + logger.info(f" Dataset repeat: {dataset_repeat}") + logger.info(f" Effective dataset size: {effective_dataset_size} (unique * repeat)") + logger.info(f" Number of epochs: {num_epochs}") + logger.info(f" Number of GPUs: {num_processes}") + logger.info(f" Gradient accumulation steps: {gradient_accumulation_steps}") + logger.info(f" Per-device batch size: {per_device_train_batch_size}") + logger.info(f" Total samples to process: {effective_dataset_size * num_epochs}") + logger.info(f" Steps per GPU: ~{total_steps_per_gpu}") + logger.info(f" Total steps (global): {total_steps_global}") + logger.info("") + logger.info(f" ✓ Each GPU will process ~{total_steps_per_gpu} steps") + logger.info(f" ✓ This ensures traversal of all {effective_dataset_size} samples") + logger.info(f" ({dataset_size} unique samples × {dataset_repeat} repeats)") + logger.info(f" ✓ Over {num_epochs} epoch(s)") + logger.info("="*80) + + step = resume_step_count + traj_loss = 0.0 + if resume_step_count > 0: + adaptation_steps = max(200, resume_step_count // 100) + spike_detection_start_step = resume_step_count + adaptation_steps + logger.info(f"Resuming from step {resume_step_count}, spike detection will start at step {spike_detection_start_step} (after {adaptation_steps} adaptation steps)") + else: + spike_detection_start_step = 100 + + for epoch_id in range(num_epochs): + epoch_seed = seed + epoch_id + torch.manual_seed(epoch_seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed(epoch_seed) + torch.cuda.manual_seed_all(epoch_seed) + + if resume_step_count > 0 and epoch_id == 0: + estimated_skip = resume_step_count // gradient_accumulation_steps + if estimated_skip > 0: + logger.info(f"Skipping {estimated_skip} data samples to resume from step {resume_step_count}...") + dataloader_iter = iter(dataloader) + for _ in tqdm(range(estimated_skip), desc="Skipping data", unit="samples", leave=False): + try: + next(dataloader_iter) + except StopIteration: + break + dataloader = dataloader_iter + logger.info(f"Successfully skipped {estimated_skip} data samples, resuming training...") + + # Track consecutive None data to detect if we're stuck in a loop + consecutive_none_count = 0 + max_consecutive_none = 100 # If we get 100 consecutive None values, something is wrong + + progress_total = int(progress_total_steps) + if progress_total <= 0: + progress_total = len(dataloader) + progress_bar = tqdm( + dataloader, + total=progress_total, + initial=resume_step_count if progress_total_steps else 0, + desc="Training steps", + unit="step", + ) + for data_idx, data in enumerate(progress_bar): + # Handle None data (can happen if all files in batch fail to load) + if data is None: + consecutive_none_count += 1 + if consecutive_none_count >= max_consecutive_none: + logger.error(f"Received {max_consecutive_none} consecutive None data samples. This suggests a serious dataset issue. Stopping training.") + raise ValueError(f"Too many consecutive None data samples ({max_consecutive_none}). Check dataset files.") + + # Log warning but continue (will skip this step) + if consecutive_none_count <= 10 or consecutive_none_count % 10 == 0: + logger.warning(f"Received None data at index {data_idx} (consecutive: {consecutive_none_count}). This may indicate missing or corrupted files. Skipping...") + + # Still increment step to keep step_count synchronized + step += 1 + dummy_loss = torch.tensor(0.0, device=accelerator.device, requires_grad=False) + model_logger.on_step_end(dummy_loss, accelerator, model, current_batch=samples) + continue + + # Reset consecutive None counter when we get valid data + consecutive_none_count = 0 + + # Normalize to list of samples for batch processing (per_device_train_batch_size > 1) + samples = data if isinstance(data, list) else [data] + + # Simplified context-based retrieval OR replay/prev_chunk_tail (aligned with multichunk eval) + context_retrieval_success = True # Set False if any sample fails (for strict mode) + _umodel = accelerator.unwrap_model(model) + _cm_frames = int(_umodel.context_memory_frames) + _cs = context_source.strip().lower() + if _cs not in ("fov", "replay", "prev_chunk_tail"): + _cs = "fov" + + if _cs == "replay" and dataset_base_path: + from src.model_training.multichunk_sample_utils import ( + replay_context_actions_from_segment_actions, + replay_context_global_indices, + synthetic_replay_context_from_segment, + ) + for d in samples: + vf = d.get("video") or [] + n_seg = min(int(prev_chunk_frames), len(vf)) if vf else 0 + ctx_pil = synthetic_replay_context_from_segment(vf, n_seg, _cm_frames) if n_seg > 0 else None + if not ctx_pil: + context_retrieval_success = False + break + d["context_frames"] = ctx_pil + d["context_source"] = "replay_synthetic" + acts = d.get("actions") + if isinstance(acts, list) and len(acts) >= n_seg: + ra = replay_context_actions_from_segment_actions(acts[:n_seg], n_seg, _cm_frames) + if ra is not None: + d["context_actions"] = ra + sf = int(d.get("start_frame", 0) or 0) + idxs = replay_context_global_indices(n_seg, _cm_frames) + d["context_frame_indices"] = [sf + int(i) for i in idxs] + + elif _cs == "prev_chunk_tail" and dataset_base_path: + from src.model_training.multichunk_sample_utils import load_prev_chunk_tail_from_disk, load_prev_chunk_tail_rt_actions + _ctx_pos = os.environ.get("CONTEXT_POSITION", "suffix").strip().lower() + _nearest_first = (_ctx_pos == "suffix") + for d in samples: + sf = int(d.get("start_frame", 0) or 0) + vn = d.get("video_name", "") + pil_list, idxs = load_prev_chunk_tail_from_disk( + dataset_base_path, str(vn), sf, _cm_frames, nearest_first=_nearest_first + ) + if not pil_list: + context_retrieval_success = False + break + d["context_frames"] = pil_list + d["context_frame_indices"] = list(idxs) if idxs else [] + d["context_source"] = "prev_chunk_tail" + ra, _ = load_prev_chunk_tail_rt_actions( + dataset_base_path, + str(vn), + sf, + _cm_frames, + use_rt_relative=use_rt_relative, + nearest_first=_nearest_first, + ) + if ra: + d["context_actions"] = ra + + elif enable_fov_retrieval and dataset_base_path: + for d in samples: + if retrieval_method == "latent_sim": + ( + context_frames, + context_actions, + context_indices, + ref_frame_idx, + video_name, + source, + ) = retrieve_context_frames_advanced( + data=d, + dataset_base_path=dataset_base_path, + top_k=fov_top_k, + drop_overlap_probability=0.1, + use_rt_relative=use_rt_relative, + retrieval_method="latent_sim", + latent_retrieval_dir=latent_retrieval_dir, + strict_overlap_labels=strict_overlap_context, + ) + else: + ( + context_frames, + context_actions, + context_indices, + ref_frame_idx, + video_name, + source, + ) = retrieve_fov_context_frames( + data=d, + dataset_base_path=dataset_base_path, + fov_retriever=fov_retriever, # unused in simplified retrieval, kept for compat + top_k=fov_top_k, # fov_top_k is number of overlap frames (4), GT frame 0 will be added automatically + use_precomputed_overlaps=True, + strict_overlap_labels=strict_overlap_context, + allow_realtime_fallback=(not strict_overlap_context), + allow_segment_fallback=(not strict_overlap_context), + ) + + if context_frames and len(context_frames) > 0: + # Use retrieved frames as context + d["context_frames"] = context_frames + if context_actions: + d["context_actions"] = context_actions + # Store retrieval metadata for visualization/debugging + d["context_frame_indices"] = context_indices + d["context_ref_frame_idx"] = ref_frame_idx + d["context_video_name"] = video_name + d["context_source"] = source + else: + context_retrieval_success = False + break + + # Strict mode: if we require context but retrieval failed, skip this step + _need_ctx_strict = ( + strict_overlap_context + and (not context_retrieval_success) + and ( + enable_fov_retrieval + or context_source.strip().lower() in ("replay", "prev_chunk_tail") + ) + ) + if _need_ctx_strict: + if step % 50 == 0 and accelerator.is_main_process: + logger.warning(f"[CONTEXT][STRICT] No context at step={step}, skipping this training sample.") + step += 1 + dummy_loss = torch.tensor(0.0, device=accelerator.device, requires_grad=False) + model_logger.on_step_end(dummy_loss, accelerator, model, current_batch=samples) + continue + + # Diagnostic: cache the raw dataset video + a VAE encode→decode roundtrip + # every CACHE_INPUT_VIDEO_EVERY steps (env, default 0=off) to check whether + # the videos the model trains on have color issues (inversion / R-B swap). + # Main-process only; no_grad side-effect — never affects the forward/grads. + _cache_every = int(os.environ.get("CACHE_INPUT_VIDEO_EVERY", "0") or "0") + if _cache_every > 0 and accelerator.is_main_process and (step % _cache_every == 0): + try: + _dbg_dir = os.path.join(model_logger.output_path, "input_video_debug") + _umodel.dump_input_video_debug(samples, step, _dbg_dir) + except Exception as _e: # noqa: BLE001 - diagnostic only + print(f"[input-video-debug] step={step} skipped: {_e}", flush=True) + + with accelerator.accumulate(model): + optimizer.zero_grad() + # One forward over full batch: data is list of B dicts when per_device_train_batch_size > 1 + # Main loss on current batch + loss = model(data) + + step += 1 + if traj_loss == 0.0: + traj_loss = loss.item() + else: + alpha = 0.01 + traj_loss = (1 - alpha) * traj_loss + alpha * loss.item() + + if step >= spike_detection_start_step and traj_loss > 0: + relative_loss = loss.item() / traj_loss + if resume_step_count > 0 and step < resume_step_count + 500: + effective_threshold = spike_threshold * 1.5 + else: + effective_threshold = spike_threshold + + should_skip = relative_loss > effective_threshold + # Keep the skip decision identical across ranks to avoid DDP hangs. + skip_t = torch.tensor(1.0 if should_skip else 0.0, device=accelerator.device, dtype=torch.float32) + if accelerator.num_processes > 1: + dist.all_reduce(skip_t, op=dist.ReduceOp.MAX) + skip_global = skip_t.item() > 0.5 + + if skip_global: + if accelerator.is_main_process: + logger.warning(f"Spike detected at step {step} (loss={loss.item():.4f}, traj_loss={traj_loss:.4f}, ratio={relative_loss:.2f}), sync skip across all ranks") + dummy_loss = torch.tensor(0.0, device=accelerator.device, requires_grad=False) + model_logger.on_step_end(dummy_loss, accelerator, model, current_batch=samples) + del loss + torch.cuda.empty_cache() + continue + accelerator.backward(loss) + loss_v = float(loss.detach().float().item()) + # CGLA SSE load-balancing aux loss (only when + # cgla_aux_loss_weight > 0). Each CGLATransformerBlock stashes + # its aux_loss during the forward; sum across blocks (mirrors + # WanVideoPipeline.training_loss). Reported for monitoring — it + # is already added to `loss` inside the pipeline. + aux_v = float("nan") + _aux_w = 0.0 + _umodel = accelerator.unwrap_model(model) + _pipe = getattr(_umodel, "pipe", None) + if _pipe is not None: + _aux_w = float(getattr(_pipe, "cgla_aux_loss_weight", 0.0) or 0.0) + if _aux_w > 0: + _aux_sum = None + for _b in getattr(getattr(_pipe, "dit", None), "blocks", []): + _a = getattr(_b, "_aux_loss", None) + if _a is not None and torch.is_tensor(_a) and torch.isfinite(_a).all(): + _aux_sum = _a if _aux_sum is None else _aux_sum + _a + if _aux_sum is not None: + aux_v = float(_aux_sum.detach().float().item()) + # grad_norm / clip only on the accumulation boundary (sync step): + # avoid a wasteful all-reduce every micro-step, and measure the + # full accumulated gradient rather than a partial one. + grad_norm = float("nan") + cur_lr = float(optimizer.param_groups[0].get("lr", 0.0)) + if accelerator.sync_gradients: + if max_grad_norm and max_grad_norm > 0: + grad_norm_t = accelerator.clip_grad_norm_(model.parameters(), max_grad_norm) + else: + grad_norm_t = accelerator.clip_grad_norm_(model.parameters(), float("inf")) + grad_norm = float(grad_norm_t.detach().float().item()) + cur_lr = float(optimizer.param_groups[0].get("lr", 0.0)) + optimizer.step() + # Live progress bar: loss every step; grad_norm/lr on the sync step; + # aux when present. + _post = {"loss": f"{loss_v:.4f}", "traj": f"{traj_loss:.4f}"} + if not math.isnan(grad_norm): + _post["gnorm"] = f"{grad_norm:.3f}" + _post["lr"] = f"{cur_lr:.2e}" + if not math.isnan(aux_v): + _post["aux"] = f"{aux_v:.4f}" + progress_bar.set_postfix(_post) + # Periodic structured line (main process, sync step only) so + # loss/grad_norm/aux land in the tee'd log file, not just the bar. + if (accelerator.sync_gradients and accelerator.is_main_process + and (step % log_interval == 0)): + _aux_str = (f" aux={aux_v:.4f}(w={_aux_w:.3f})" + if not math.isnan(aux_v) else "") + logger.info( + f"[train] step={step} loss={loss_v:.4f} " + f"traj={traj_loss:.4f} grad_norm={grad_norm:.4f} " + f"lr={cur_lr:.2e}{_aux_str}" + ) + model_logger.on_step_end(loss, accelerator, model, current_batch=samples, + aux_loss=aux_v) + scheduler.step() + + if max_train_steps and step >= max_train_steps: + if progress_total_steps: + progress_bar.n = min(step, progress_bar.total) if progress_bar.total is not None else step + progress_bar.refresh() + if accelerator.is_main_process: + logger.info(f"[TRAIN] Reached max_train_steps={max_train_steps}; stopping without epoch checkpoint.") + accelerator.wait_for_everyone() + return + + model_logger.on_epoch_end(accelerator, model, epoch_id) diff --git a/code/src/model_training/training_modules/wan_training_module.py b/code/src/model_training/training_modules/wan_training_module.py new file mode 100644 index 0000000000000000000000000000000000000000..e7fba14908b2c26bc903638c9c1767c5d78c7a70 --- /dev/null +++ b/code/src/model_training/training_modules/wan_training_module.py @@ -0,0 +1,869 @@ +import hashlib +import json +import logging +import os +import random +from typing import Any, Dict, Optional + +import torch +from safetensors.torch import load_file as safe_load_file + +from diffsynth.pipelines.wan_video_new import WanVideoPipeline, ModelConfig +from src.model_training.transformers_compat import patch_transformers_hybrid_cache + +patch_transformers_hybrid_cache() +from diffsynth.trainers.utils import DiffusionTrainingModule +from diffsynth.models.memory.geometry_spatial_memory import GeometrySpatialMemory +from diffsynth.models.memory.mixture_of_contexts import MixtureOfContexts +from diffsynth.models.memory.spatial_grid_memory import SpatialCrossAttnReadout, SpatialGridMemory +from src.model_training.fov_retrieval import flip_yaw_rt_list + +logger = logging.getLogger(__name__) + + +class WanTrainingModule(DiffusionTrainingModule): + def __init__( + self, + model_paths=None, model_id_with_origin_paths=None, + tokenizer_path=None, + trainable_models=None, + lora_base_model=None, lora_target_modules="q,k,v,o,ffn.0,ffn.2", lora_rank=32, + use_gradient_checkpointing=True, + use_gradient_checkpointing_offload=False, + extra_inputs=None, + timestep_shift=1.0, + resume_from_checkpoint=None, + dataset_base_path: Optional[str] = None, + enable_context_memory=False, + context_memory_frames=8, + training_mode="context", # "context" mode for Context Memory (inpainting) + context_drop_prob: float = 0.0, + context_drop_seed: int = 42, + omit_context_actions: bool = False, # Context-as-Memory: no context RT injection + context_noise_prob=0.0, + context_noise_std=0.02, + context_fixed_noise_std=None, # Experiment 7: Fixed noise std (e.g., 0.1) to align training-inference + teacher_forcing_prob=0.0, + yaw_flip_aug: bool = False, # 50% prob flip yaw (ACTION_FOLLOWING direction sensitivity) + context_per_frame_vae: bool = False, # Encode each context frame separately (1 latent per raw frame) + context_source: str = "fov", # fov | replay | prev_chunk_tail (multichunk-aligned context construction) + use_framepack_memory: bool = False, + context_temporal_decay: float = 1.0, + context_attention_weight: float = 1.0, + use_framepack_length_compress: bool = False, + framepack_ratio: int = 2, + framepack_length_strategy: str = "distance_merge", + framepack_recent_keep_ratio: float = 0.5, + framepack_multiscale_w2: float = 0.25, + framepack_multiscale_w4: float = 0.15, + use_spatial_memory: bool = False, + use_spatial_memory_legacy: bool = False, + spatial_memory_tokens: int = 64, + spatial_memory_grid: int = 8, + spatial_memory_inject_mode: str = "concat_text", + use_geometry_spatial_memory: bool = False, + geometry_spatial_memory_tokens: int = 64, + geometry_spatial_memory_grid: int = 8, + geometry_spatial_memory_temporal_bins: int = 4, + geometry_spatial_memory_inject_mode: str = "concat_text", + use_moc: bool = False, + moc_temperature: float = 1.0, + moc_top_k: int = 0, + # Note: Self-forcing parameters removed - using standard training only + ): + super().__init__() + # Load models + model_configs = [] + if model_paths is not None: + model_paths = json.loads(model_paths) + model_configs += [ModelConfig(path=path) for path in model_paths] + if model_id_with_origin_paths is not None: + model_id_with_origin_paths = model_id_with_origin_paths.split(",") + model_configs += [ModelConfig(model_id=i.split(":")[0], origin_file_pattern=i.split(":")[1]) for i in model_id_with_origin_paths] + from_pretrained_kw = {"torch_dtype": torch.bfloat16, "device": "cpu", "model_configs": model_configs} + if tokenizer_path: + from_pretrained_kw["tokenizer_config"] = ModelConfig(path=tokenizer_path) + self.pipe = WanVideoPipeline.from_pretrained(**from_pretrained_kw) + + # Store timestep_shift for later use (e.g., after video sampling) + self.timestep_shift = timestep_shift + + # Reset training scheduler + self.pipe.scheduler.set_timesteps(1000, training=True, shift=timestep_shift) + + # Freeze untrainable models + self.pipe.freeze_except([] if trainable_models is None else trainable_models.split(",")) + + # Add LoRA to the base models + if lora_base_model is not None: + model = self.add_lora_to_model( + getattr(self.pipe, lora_base_model), + target_modules=lora_target_modules.split(","), + lora_rank=lora_rank + ) + setattr(self.pipe, lora_base_model, model) + + # Load checkpoint if provided + if resume_from_checkpoint is not None: + logger.info(f"Loading LoRA checkpoint from: {resume_from_checkpoint}") + if not os.path.exists(resume_from_checkpoint): + raise FileNotFoundError(f"Checkpoint file not found: {resume_from_checkpoint}") + checkpoint_state_dict = safe_load_file(resume_from_checkpoint) + logger.info(f"Checkpoint contains {len(checkpoint_state_dict)} parameters") + # The checkpoint was saved with remove_prefix_in_ckpt, so keys don't have the prefix + # The model (pipe.dit) state_dict keys also don't have the prefix, so they should match + # Use strict=False to allow partial loading + missing_keys, unexpected_keys = model.load_state_dict(checkpoint_state_dict, strict=False) + if missing_keys: + logger.warning(f"{len(missing_keys)} keys were missing when loading checkpoint") + if len(missing_keys) <= 10: + logger.debug(f"Missing keys: {missing_keys}") + if unexpected_keys: + logger.warning(f"{len(unexpected_keys)} unexpected keys in checkpoint (will be ignored)") + if len(unexpected_keys) <= 10: + logger.debug(f"Unexpected keys: {unexpected_keys}") + loaded_count = len(checkpoint_state_dict) - len(missing_keys) - len(unexpected_keys) + logger.info(f"Successfully loaded {loaded_count} parameters from checkpoint!") + + # Store other configs + self.use_gradient_checkpointing = use_gradient_checkpointing + self.use_gradient_checkpointing_offload = use_gradient_checkpointing_offload + self.extra_inputs = extra_inputs.split(",") if extra_inputs is not None else [] + self.dataset_base_path = dataset_base_path + + # Context Memory (Context as Memory) configuration + self.enable_context_memory = enable_context_memory + self.context_memory_frames = context_memory_frames + self.training_mode = training_mode # "predict", "context", or "condition" + self.context_drop_prob = float(context_drop_prob or 0.0) + self.context_drop_seed = int(context_drop_seed or 42) + self.omit_context_actions = bool(omit_context_actions) + self.context_per_frame_vae = bool(context_per_frame_vae) + self.context_source = (context_source or "fov").strip().lower() + if self.context_source not in ("fov", "replay", "prev_chunk_tail"): + self.context_source = "fov" + self.context_noise_prob = context_noise_prob + self.context_noise_std = context_noise_std + self.context_fixed_noise_std = context_fixed_noise_std # Experiment 7: Fixed noise for training-inference alignment + self.teacher_forcing_prob = teacher_forcing_prob + self.teacher_forcing_enabled = teacher_forcing_prob > 0.0 + self.yaw_flip_aug = bool(yaw_flip_aug) + # Memory baselines runtime flags (train + sampling path shared). + self.use_framepack_memory = bool(use_framepack_memory) + self.context_temporal_decay = float(context_temporal_decay or 1.0) + self.context_attention_weight = float(context_attention_weight or 1.0) + self.use_framepack_length_compress = bool(use_framepack_length_compress) + self.framepack_ratio = int(framepack_ratio or 2) + self.framepack_length_strategy = str(framepack_length_strategy or "distance_merge").lower() + self.framepack_recent_keep_ratio = float(framepack_recent_keep_ratio or 0.5) + self.framepack_multiscale_w2 = float(framepack_multiscale_w2 or 0.25) + self.framepack_multiscale_w4 = float(framepack_multiscale_w4 or 0.15) + # Mirror key flags to pipe for inference-time sampling monitor. + self.pipe.use_framepack_memory = self.use_framepack_memory + self.pipe.context_temporal_decay = self.context_temporal_decay + self.pipe.context_attention_weight = self.context_attention_weight + self.pipe.use_framepack_length_compress = self.use_framepack_length_compress + self.pipe.framepack_ratio = self.framepack_ratio + self.pipe.framepack_length_strategy = self.framepack_length_strategy + self.pipe.framepack_recent_keep_ratio = self.framepack_recent_keep_ratio + self.pipe.framepack_multiscale_w2 = self.framepack_multiscale_w2 + self.pipe.framepack_multiscale_w4 = self.framepack_multiscale_w4 + self.use_moc = bool(use_moc) + self.moc_temperature = float(moc_temperature or 1.0) + self.moc_top_k = int(moc_top_k or 0) + self.moc_module = MixtureOfContexts( + temperature=self.moc_temperature, + top_k=self.moc_top_k, + ) if self.use_moc else None + self.pipe.use_moc = self.use_moc + self.pipe.moc_module = self.moc_module + self.pipe.use_spatial_memory = bool(use_spatial_memory) + self.pipe.use_spatial_memory_legacy = bool(use_spatial_memory_legacy) + self.pipe.spatial_memory_tokens = int(spatial_memory_tokens or 64) + self.pipe.spatial_memory_inject_mode = str(spatial_memory_inject_mode or "concat_text") + self.spatial_memory_module = None + self.spatial_memory_readout_module = None + if self.pipe.use_spatial_memory and not self.pipe.use_spatial_memory_legacy: + dim = int(getattr(self.pipe.dit, "dim")) + grid_size = int(spatial_memory_grid or 8) + self.pipe.spatial_memory_grid = grid_size + self.spatial_memory_module = SpatialGridMemory( + dim=dim, + grid_size=grid_size, + num_tokens=self.pipe.spatial_memory_tokens, + ) + self.pipe.spatial_memory_module = self.spatial_memory_module + if self.pipe.spatial_memory_inject_mode == "cross_attn_readout": + self.spatial_memory_readout_module = SpatialCrossAttnReadout(dim=dim, num_heads=8) + self.pipe.spatial_memory_readout_module = self.spatial_memory_readout_module + else: + self.pipe.spatial_memory_module = None + self.pipe.spatial_memory_readout_module = None + self.use_geometry_spatial_memory = bool(use_geometry_spatial_memory) + self.geometry_spatial_memory_module = None + self.geometry_spatial_memory_readout_module = None + self.pipe.use_geometry_spatial_memory = self.use_geometry_spatial_memory + self.pipe.geometry_spatial_memory_inject_mode = str( + geometry_spatial_memory_inject_mode or "concat_text" + ) + if self.use_geometry_spatial_memory: + dim = int(getattr(self.pipe.dit, "dim")) + self.geometry_spatial_memory_module = GeometrySpatialMemory( + dim=dim, + latent_channels=int(getattr(self.pipe.dit, "in_dim", 16)), + patch_size=tuple(getattr(self.pipe.dit, "patch_size", (1, 2, 2))), + grid_size=int(geometry_spatial_memory_grid or 8), + temporal_bins=int(geometry_spatial_memory_temporal_bins or 4), + num_tokens=int(geometry_spatial_memory_tokens or 64), + ) + self.geometry_spatial_memory_module.initialize_from_dit_patch_embedding( + self.pipe.dit.patch_embedding + ) + dit_parameter = next(self.pipe.dit.parameters()) + self.geometry_spatial_memory_module = self.geometry_spatial_memory_module.to( + device=dit_parameter.device, + dtype=dit_parameter.dtype, + ) + self.pipe.geometry_spatial_memory_module = self.geometry_spatial_memory_module + if self.pipe.geometry_spatial_memory_inject_mode == "cross_attn_readout": + self.geometry_spatial_memory_readout_module = SpatialCrossAttnReadout( + dim=dim, + num_heads=8, + ).to(device=dit_parameter.device, dtype=dit_parameter.dtype) + self.pipe.geometry_spatial_memory_readout_module = ( + self.geometry_spatial_memory_readout_module + ) + else: + self.pipe.geometry_spatial_memory_module = None + self.pipe.geometry_spatial_memory_readout_module = None + # Note: Self-forcing removed - using standard training only + self.current_step = 0 # Track current training step (for logging/debugging) + + def _forward_preprocess_batch(self, samples: list) -> dict: + """Batch preprocessing for Stage 1 Interactive (no context). data is list of sample dicts.""" + if not samples: + raise ValueError("samples cannot be empty in _forward_preprocess_batch") + batch_size = len(samples) + prompts = [] + video_frames_list = [] + actions_list = [] + for s in samples: + p = s.get("prompt") + if p is None: + raise ValueError("sample['prompt'] is missing or None") + prompts.append(str(p) if not isinstance(p, str) else p) + video_frames_list.append(s["video"]) + if "actions" in s and s["actions"] is not None: + acts = s["actions"] + if getattr(self, 'yaw_flip_aug', False) and isinstance(acts, list) and len(acts) > 0 and isinstance(acts[0], (list, tuple)) and len(acts[0]) >= 12 and random.random() < 0.5: + acts = flip_yaw_rt_list(acts) + if isinstance(acts, torch.Tensor): + actions_list.append(acts) + elif isinstance(acts, list) and len(acts) > 0: + actions_list.append(torch.tensor(acts, dtype=torch.float32)) + else: + actions_list.append(None) + else: + actions_list.append(None) + + # input_video: list of lists (each inner list = PIL images for one video) + input_video = video_frames_list + first = samples[0] + h, w = first["video"][0].size[1], first["video"][0].size[0] + num_frames = len(first["video"]) + + inputs_posi = {"prompt": prompts} + inputs_nega = {} + inputs_shared = { + "input_video": input_video, + "height": h, + "width": w, + "num_frames": num_frames, + "batch_size": batch_size, + "cfg_scale": 1, + "tiled": False, + "rand_device": self.pipe.device, + "use_gradient_checkpointing": self.use_gradient_checkpointing, + "use_gradient_checkpointing_offload": self.use_gradient_checkpointing_offload, + "cfg_merge": False, + "vace_scale": 1, + } + + ref_action = next((a for a in actions_list if a is not None), None) + if ref_action is not None and batch_size == 1: + inputs_shared["actions"] = ref_action.detach().cpu().tolist() if isinstance(ref_action, torch.Tensor) else ref_action + elif ref_action is not None: + device = self.pipe.device + dtype = ref_action.dtype + stacked = [] + for a in actions_list: + if a is not None: + stacked.append(a.to(device=device)) + else: + stacked.append(torch.zeros_like(ref_action, device=device, dtype=dtype)) + inputs_shared["actions"] = torch.stack(stacked) + else: + inputs_shared["actions"] = None + + for unit in self.pipe.units: + inputs_shared, inputs_posi, inputs_nega = self.pipe.unit_runner(unit, self.pipe, inputs_shared, inputs_posi, inputs_nega) + return {**inputs_shared, **inputs_posi} + + def _build_context_with_anchor(self, context_frames, context_actions=None, expected_k=None): + """Training-side anchor helper: keep last frame as mandatory anchor and keep action length aligned.""" + frames = list(context_frames or []) + actions = list(context_actions or []) if context_actions is not None else [] + if not frames or not getattr(self, "use_anchor_frame", False): + return frames, actions + k = int(expected_k) if (expected_k is not None and int(expected_k) > 0) else len(frames) + if len(frames) > k: + frames = frames[-k:] + if actions: + actions = actions[-k:] + if actions: + if len(actions) < len(frames): + actions = actions + [actions[-1]] * (len(frames) - len(actions)) + elif len(actions) > len(frames): + actions = actions[:len(frames)] + return frames, actions + + def _forward_preprocess_batch_context(self, samples: list) -> dict: + """Batch preprocessing for Stage 2 Context Memory. Batch-level drop: if drop, all samples get no context.""" + if not samples: + raise ValueError("samples cannot be empty in _forward_preprocess_batch_context") + batch_size = len(samples) + first = samples[0] + + def _should_drop_context(_data) -> bool: + p = float(getattr(self, "context_drop_prob", 0.0) or 0.0) + if p <= 0.0: + return False + if p >= 1.0: + return True + vn = str(_data.get("video_name", "")) + sf = str(_data.get("start_frame", "")) + key = f"{int(getattr(self, 'context_drop_seed', 42))}|{vn}|{sf}" + h = hashlib.md5(key.encode("utf-8")).hexdigest() + u = int(h[:8], 16) / 0xFFFFFFFF + return u < p + + # Batch-level drop: use first sample to decide for whole batch + dropped_context = _should_drop_context(first) + # IMPORTANT (DDP safety): ensure all ranks make the same drop decision. + # If some ranks drop context while others keep it, modules conditioned on context + # (e.g. implicit encoder / compressor) become unused on a subset of ranks and can + # deadlock gradient sync / trigger NCCL watchdog timeouts. + try: + import torch.distributed as dist + if dist.is_available() and dist.is_initialized(): + flag = torch.tensor([1 if dropped_context else 0], device=self.pipe.device, dtype=torch.int64) + dist.broadcast(flag, src=0) + dropped_context = bool(int(flag.item())) + except Exception: + pass + + prompts = [] + video_frames_list = [] + actions_list = [] + context_latents_list = [] + context_actions_list = [] + geometry_memory_latents_list = [] + expected_k = self.context_memory_frames + training_mode = getattr(self, 'training_mode', 'context') + + target_h = first["video"][0].size[1] + target_w = first["video"][0].size[0] + num_frames = len(first["video"]) + + from PIL import Image + + for s in samples: + p = s.get("prompt") + if p is None: + raise ValueError("sample['prompt'] is missing or None") + prompts.append(str(p) if not isinstance(p, str) else p) + video_frames_list.append(s["video"]) + + if "actions" in s and s["actions"] is not None: + acts = s["actions"] + if getattr(self, 'yaw_flip_aug', False) and isinstance(acts, list) and len(acts) > 0 and isinstance(acts[0], (list, tuple)) and len(acts[0]) >= 12 and random.random() < 0.5: + acts = flip_yaw_rt_list(acts) + if isinstance(acts, torch.Tensor): + actions_list.append(acts) + elif isinstance(acts, list) and len(acts) > 0: + actions_list.append(torch.tensor(acts, dtype=torch.float32)) + else: + actions_list.append(None) + else: + actions_list.append(None) + + geometry_frames = s.get("geometry_memory_frames") or [] + if self.use_geometry_spatial_memory: + if not geometry_frames: + raise ValueError( + "Geometry-grounded Spatial Memory requires sample['geometry_memory_frames']. " + "Provide TSDF/point-cloud renders through the configured metadata column." + ) + resized_geometry = [] + for frame in geometry_frames: + if hasattr(frame, "resize") and hasattr(frame, "size"): + gw, gh = frame.size + if gh != target_h or gw != target_w: + frame = frame.resize((target_w, target_h), Image.Resampling.LANCZOS) + resized_geometry.append(frame) + with torch.no_grad(): + geometry_video = self.pipe.preprocess_video(resized_geometry) + if geometry_video.dim() == 4: + geometry_video = geometry_video.unsqueeze(0) + geometry_latents = self.pipe.vae.encode( + [geometry_video[i] for i in range(geometry_video.shape[0])], + device=self.pipe.device, + tiled=False, + tile_size=None, + tile_stride=None, + ) + geometry_memory_latents_list.append( + geometry_latents.to(dtype=self.pipe.torch_dtype, device=self.pipe.device) + ) + else: + geometry_memory_latents_list.append(None) + + if dropped_context: + context_latents_list.append(None) + context_actions_list.append(None) + continue + + ctx_frames = s.get("context_frames") or [] + ctx_actions = [] if getattr(self, "omit_context_actions", False) else (s.get("context_actions") or []) # ctx=1: no context action + context_indices = s.get("context_frame_indices", []) + start_frame = s.get("start_frame", None) + end_frame = s.get("end_frame", None) + + if ctx_frames and context_indices and start_frame is not None and end_frame is not None: + filtered_frames, filtered_actions = [ctx_frames[0]], [] + if ctx_actions: + filtered_actions.append(ctx_actions[0]) + for i in range(1, len(ctx_frames)): + idx = context_indices[i] if i < len(context_indices) else None + if idx is None or idx < start_frame or idx > end_frame: + filtered_frames.append(ctx_frames[i]) + if ctx_actions and i < len(ctx_actions): + filtered_actions.append(ctx_actions[i]) + ctx_frames, ctx_actions = filtered_frames, filtered_actions if filtered_actions else ctx_actions + + if not ctx_frames and len(s["video"]) > expected_k: + ctx_frames = s["video"][:expected_k] + if s.get("actions") and len(s["actions"]) >= expected_k: + ctx_actions = s["actions"][:expected_k] + + if not ctx_frames: + context_latents_list.append(None) + context_actions_list.append(None) + continue + + resized = [] + for f in ctx_frames: + if hasattr(f, 'resize') and hasattr(f, 'size'): + w, h = f.size + if h != target_h or w != target_w: + f = f.resize((target_w, target_h), Image.Resampling.LANCZOS) + resized.append(f) + ctx_frames = resized + + if len(ctx_frames) < expected_k: + last = ctx_frames[-1] if ctx_frames else Image.new('RGB', (target_w, target_h), (0, 0, 0)) + ctx_frames = ctx_frames + [last] * (expected_k - len(ctx_frames)) + if ctx_actions: + ctx_actions = ctx_actions + [ctx_actions[-1]] * (expected_k - len(ctx_actions)) + elif len(ctx_frames) > expected_k: + ctx_frames = ctx_frames[:expected_k] + ctx_actions = ctx_actions[:expected_k] if ctx_actions else [] + + ctx_frames, ctx_actions = self._build_context_with_anchor( + ctx_frames, + context_actions=ctx_actions, + expected_k=expected_k, + ) + + with torch.no_grad(): + if getattr(self, "context_per_frame_vae", False): + # Each context frame -> 1 latent token (no temporal downsample); context_actions remain one per raw frame + context_latents_per_sample = [] + for f in ctx_frames: + frame_video = self.pipe.preprocess_video([f]) # (1, C, 1, H, W) + frame_sq = frame_video.squeeze(0) # (C, 1, H, W) + lat_one = self.pipe.vae.encode([frame_sq], device=self.pipe.device, tiled=False, tile_size=None, tile_stride=None) + context_latents_per_sample.append(lat_one) + lat = torch.cat(context_latents_per_sample, dim=2) # (1, C, K, H//8, W//8) + else: + ctx_video = self.pipe.preprocess_video(ctx_frames) + if ctx_video.dim() == 4: + ctx_video = ctx_video.unsqueeze(0) + lat = self.pipe.vae.encode([ctx_video[i] for i in range(ctx_video.shape[0])], device=self.pipe.device, tiled=False, tile_size=None, tile_stride=None) + context_latents_list.append(lat.to(dtype=self.pipe.torch_dtype, device=self.pipe.device)) + + if ctx_actions: + if isinstance(ctx_actions[0], (list, tuple)): + context_actions_list.append(torch.tensor(ctx_actions, dtype=torch.float32)) + else: + context_actions_list.append(torch.tensor(ctx_actions, dtype=torch.float32)) + else: + context_actions_list.append(None) + + input_video = video_frames_list + inputs_posi = {"prompt": prompts} + inputs_nega = {} + inputs_shared = { + "input_video": input_video, + "height": target_h, + "width": target_w, + "num_frames": num_frames, + "batch_size": batch_size, + "cfg_scale": 1, + "tiled": False, + "rand_device": self.pipe.device, + "use_gradient_checkpointing": self.use_gradient_checkpointing, + "use_gradient_checkpointing_offload": self.use_gradient_checkpointing_offload, + "cfg_merge": False, + "vace_scale": 1, + } + + # DDP safety: ensure *all* ranks either have context (and thus use context-conditioned modules) + # or all ranks drop it. Using an all-reduce MIN means if any rank lacks context, we drop globally. + has_context_step = (not dropped_context) and any(x is not None for x in context_latents_list) + try: + import torch.distributed as dist + if dist.is_available() and dist.is_initialized(): + flag = torch.tensor([1 if has_context_step else 0], device=self.pipe.device, dtype=torch.int64) + dist.all_reduce(flag, op=dist.ReduceOp.MIN) + has_context_step = bool(int(flag.item())) + except Exception: + pass + if not has_context_step: + dropped_context = True + + if not dropped_context and any(x is not None for x in context_latents_list): + valid = [x for x in context_latents_list if x is not None] + if valid: + ref = valid[0] + device, dtype = self.pipe.device, ref.dtype + stacked_ctx = [] + for x in context_latents_list: + if x is not None: + stacked_ctx.append(x.to(device=device)) + else: + stacked_ctx.append(torch.zeros_like(ref, device=device, dtype=dtype)) + inputs_shared["context_latents"] = torch.cat(stacked_ctx, dim=0) + inputs_shared["num_context_frames"] = ref.shape[2] + inputs_shared["training_mode"] = training_mode + inputs_shared["context_noise_prob"] = getattr(self, 'context_noise_prob', 0.0) + inputs_shared["context_noise_std"] = getattr(self, 'context_noise_std', 0.02) + if self.context_fixed_noise_std is not None: + inputs_shared["context_fixed_noise_std"] = self.context_fixed_noise_std + inputs_shared["context_position"] = os.environ.get("CONTEXT_POSITION", "suffix") + inputs_shared["omit_context_actions"] = getattr(self, "omit_context_actions", False) + inputs_shared["context_attention_weight"] = getattr(self, "context_attention_weight", 1.0) + inputs_shared["use_anchor_frame"] = getattr(self, "use_anchor_frame", False) + inputs_shared["context_temporal_decay"] = getattr(self, "context_temporal_decay", 1.0) + inputs_shared["use_spatial_memory"] = getattr(self.pipe, "use_spatial_memory", False) + inputs_shared["spatial_memory_tokens"] = int(getattr(self.pipe, "spatial_memory_tokens", 64) or 64) + inputs_shared["use_spatial_memory_legacy"] = bool(getattr(self.pipe, "use_spatial_memory_legacy", False)) + inputs_shared["spatial_memory_module"] = getattr(self.pipe, "spatial_memory_module", None) + inputs_shared["spatial_memory_inject_mode"] = getattr(self.pipe, "spatial_memory_inject_mode", "concat_text") + inputs_shared["spatial_memory_readout_module"] = getattr(self.pipe, "spatial_memory_readout_module", None) + inputs_shared["use_framepack_memory"] = bool(getattr(self, "use_framepack_memory", False)) + if self.use_moc and self.moc_module is not None: + inputs_shared["use_moc"] = True + inputs_shared["moc_module"] = self.moc_module + nf_list = [s.get("non_fov_frames") or [] for s in samples] + if any(nf for nf in nf_list): + inputs_shared["non_fov_frames_list"] = nf_list + + if self.use_geometry_spatial_memory: + if not all(x is not None for x in geometry_memory_latents_list): + raise ValueError("Geometry memory is missing for one or more samples in the batch.") + inputs_shared["geometry_memory_latents"] = torch.cat( + geometry_memory_latents_list, + dim=0, + ) + inputs_shared["use_geometry_spatial_memory"] = True + inputs_shared["geometry_spatial_memory_module"] = ( + self.geometry_spatial_memory_module + ) + inputs_shared["geometry_spatial_memory_inject_mode"] = ( + self.pipe.geometry_spatial_memory_inject_mode + ) + inputs_shared["geometry_spatial_memory_readout_module"] = ( + self.geometry_spatial_memory_readout_module + ) + + ctx_acts_valid = [a for a in context_actions_list if a is not None] + if not getattr(self, "omit_context_actions", False) and ctx_acts_valid: + ref_act = ctx_acts_valid[0] + target_len = ref_act.shape[0] # num_context_frames (K) + stacked_ca = [] + for a in context_actions_list: + if a is not None: + a = a.to(device=device) + if a.shape[0] != target_len: + if a.shape[0] > target_len: + a = a[:target_len] + else: + pad = a.new_zeros(target_len - a.shape[0], a.shape[-1]) + a = torch.cat([a, pad], dim=0) + stacked_ca.append(a) + else: + stacked_ca.append(torch.zeros_like(ref_act, device=device, dtype=ref_act.dtype)) + inputs_shared["context_actions"] = torch.stack(stacked_ca) + + ref_action = next((a for a in actions_list if a is not None), None) + if ref_action is not None and batch_size == 1: + inputs_shared["actions"] = ref_action.detach().cpu().tolist() if isinstance(ref_action, torch.Tensor) else ref_action + elif ref_action is not None: + device = self.pipe.device + dtype = ref_action.dtype + stacked = [] + for a in actions_list: + if a is not None: + stacked.append(a.to(device=device)) + else: + stacked.append(torch.zeros_like(ref_action, device=device, dtype=dtype)) + inputs_shared["actions"] = torch.stack(stacked) + else: + inputs_shared["actions"] = None + + for unit in self.pipe.units: + inputs_shared, inputs_posi, inputs_nega = self.pipe.unit_runner(unit, self.pipe, inputs_shared, inputs_posi, inputs_nega) + return {**inputs_shared, **inputs_posi} + + @staticmethod + def _translate_condition_keys(d): + """Map VWM CamVideoDataset condition_* keys to context-memory keys.""" + if not isinstance(d, dict): + return d + if "condition_frames" in d and "context_frames" not in d: + d["context_frames"] = d.pop("condition_frames") + if "condition_actions" in d and "context_actions" not in d: + d["context_actions"] = d.pop("condition_actions") + if "condition_frame_indices" in d and "context_frame_indices" not in d: + d["context_frame_indices"] = d.pop("condition_frame_indices") + if "use_condition_context_frames" in d: + d.pop("use_condition_context_frames") + if "condition_source" in d: + d.pop("condition_source", None) + return d + + def forward_preprocess(self, data): + if data is None: + raise ValueError("data cannot be None in forward_preprocess") + samples = data if isinstance(data, list) else [data] + samples = [self._translate_condition_keys(d) for d in samples] + if self.enable_context_memory: + return self._forward_preprocess_batch_context(samples) + return self._forward_preprocess_batch(samples) + + def _ensure_input_latents(self, inputs: Dict[str, Any], *, strict: bool = False) -> Dict[str, Any]: + if "input_latents" in inputs: + return inputs + import warnings + video_obj = inputs.get("input_video", None) + if video_obj is None: + video_obj = inputs.get("video", None) + vae = getattr(self.pipe, "vae", None) + rebuild_err = None + if video_obj is not None and vae is not None and hasattr(vae, "encode"): + try: + if isinstance(video_obj, list): + video_tensor = self.pipe.preprocess_video(video_obj) + else: + video_tensor = video_obj + if hasattr(video_tensor, "dim"): + video_sq = video_tensor.squeeze(0) if video_tensor.dim() == 5 else video_tensor + with torch.no_grad(): + try: + lat = vae.encode(video_tensor, device=self.pipe.device, tiled=False, tile_size=None, tile_stride=None) + except Exception as e_first: + # Retry with the list form (matches how context latents are + # encoded in forward_preprocess). Surface BOTH errors if this + # also fails, so the real cause (often CUDA OOM right after the + # periodic sampling monitor ran) isn't hidden behind a KeyError. + try: + lat = vae.encode([video_sq], device=self.pipe.device, tiled=False, tile_size=None, tile_stride=None) + except Exception as e_retry: + raise RuntimeError( + f"VAE encode failed -- tensor form: {e_first!r}; list form: {e_retry!r}" + ) from e_retry + if isinstance(lat, (list, tuple)): + lat = lat[0] + if hasattr(lat, "dim") and lat.dim() == 4: + lat = lat.unsqueeze(0) + inputs["input_latents"] = lat.to(dtype=torch.bfloat16, device=self.pipe.device) + return inputs + except Exception as e: + rebuild_err = e + warnings.warn(f"Failed to rebuild input_latents: {e}") + msg = ( + "input_latents missing and auto-rebuild failed" + + (f" (rebuild error: {rebuild_err!r})" if rebuild_err + else " (no input_video/video or vae unavailable)") + + f". available input keys={sorted(list(inputs.keys()))}" + ) + if strict: + raise KeyError(msg) + warnings.warn(msg) + return inputs + + def restore_after_sampling(self): + """Restore training-time pipe state clobbered by the periodic sampling + monitor (``pipe.__call__``) and release its GPU cache. Called by the + ModelLogger after every paper-process sampling step so the next training + step is unaffected. + + - Scheduler: sampling runs ``set_timesteps(num_inference_steps, + training=False)``; ``training_loss`` reads ``self.scheduler.timesteps`` + directly, so we must re-apply the training schedule (1000 steps, + ``training=True``) -- otherwise every later step silently samples from + inference timesteps/sigmas (wrong loss). ``self.timestep_shift`` is + stored in ``__init__`` for exactly this. + - Cache: ~50 denoise steps + an 81-frame VAE decode leave the main rank's + GPU fragmented; releasing the cache prevents the next step's target + VAE encode (which auto-rebuilds ``input_latents``) from OOMing. + """ + self.pipe.scheduler.set_timesteps(1000, training=True, shift=self.timestep_shift) + import gc + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + def dump_input_video_debug(self, samples, step, out_dir, fps=15): + """Cache the raw dataset video + a VAE encode→decode roundtrip + color + stats to diagnose train/infer color consistency (e.g. color inversion). + + Writes, for ``samples[0]``, to ``/step_{N:07d}_*``: + * ``_raw.mp4`` — the raw input frames (``d["video"]``) + * ``_vae_roundtrip.mp4`` — ``preprocess_video → vae.encode → vae.decode`` + (mirrors the exact encode/decode calls used in training's + ``_ensure_input_latents`` and the pipeline's decode path, so a color + artifact here implicates the VAE / preprocessing, not the DiT/CGLA) + * ``_stats.json`` — per-channel mean RGB (raw vs roundtrip), + latent mean/std/min/max, and inversion / R-B-swap flags + + All under ``torch.no_grad`` and wrapped so a failure never aborts training + (returns the error string). No VRAM management in training → + ``load_models_to_device`` is a no-op, so the VAE stays on the train device. + """ + import json as _json + import os as _os + import numpy as _np + from diffsynth import save_video + + if not samples: + return "no samples" + d = samples[0] + frames = d.get("video") or [] + if not frames: + return "no video frames in sample" + vae = getattr(self.pipe, "vae", None) + if vae is None or not hasattr(vae, "encode") or not hasattr(vae, "decode"): + return "vae unavailable" + + _os.makedirs(out_dir, exist_ok=True) + tag = f"step_{int(step):07d}" + vn = str(d.get("video_name", "")) + sf = int(d.get("start_frame", 0) or 0) + + raw_path = _os.path.join(out_dir, f"{tag}_raw.mp4") + roundtrip_path = _os.path.join(out_dir, f"{tag}_vae_roundtrip.mp4") + stats_path = _os.path.join(out_dir, f"{tag}_stats.json") + + # Save raw input frames as-is (these are the dataset frames the model sees). + save_video(list(frames), raw_path, fps=fps, quality=5) + + stats = {"step": int(step), "video_name": vn, "start_frame": sf, + "num_frames": len(frames), "vae_dtype": str(next(vae.parameters()).dtype)} + + # VAE encode → decode roundtrip (same calls as training's input-latent + # rebuild + the pipeline decode). Cast frames through preprocess_video so + # the normalization (min/max=-1/1) matches what the model trains on. + rec_frames = None + try: + video_tensor = self.pipe.preprocess_video(list(frames)) # (1,C,T,H,W) in [-1,1] + with torch.no_grad(): + lat = vae.encode(video_tensor, device=self.pipe.device, tiled=False, + tile_size=None, tile_stride=None) + if isinstance(lat, (list, tuple)): + lat = lat[0] + rec = vae.decode(lat, device=self.pipe.device, tiled=False, + tile_size=None, tile_stride=None) + rec_frames = self.pipe.vae_output_to_video(rec) # list[PIL] + save_video(list(rec_frames), roundtrip_path, fps=fps, quality=5) + except Exception as e: # noqa: BLE001 - diagnostic only + stats["roundtrip_error"] = repr(e) + + # Color stats: detect inversion (rec ≈ 255 - raw) and R/B swap. + def _mean_rgb(pil_list): + arr = _np.stack([_np.asarray(f.convert("RGB"), dtype=_np.float32) for f in pil_list]) + return arr.reshape(-1, 3).mean(0).tolist() + + raw_mean = _mean_rgb(frames) + stats["raw_mean_rgb"] = raw_mean + if rec_frames is not None: + rec_mean = _mean_rgb(rec_frames) + stats["roundtrip_mean_rgb"] = rec_mean + r_raw, g_raw, b_raw = raw_mean + r_rec, g_rec, b_rec = rec_mean + stats["color_inversion_flag"] = bool( + abs((255 - r_raw) - r_rec) < abs(r_raw - r_rec) or + abs((255 - g_raw) - g_rec) < abs(g_raw - g_rec) + ) + stats["rb_swap_flag"] = bool(abs(r_raw - b_rec) < abs(r_raw - r_rec)) + if torch.is_tensor(lat): + _lat = lat.detach().float() + stats["latent_mean"] = float(_lat.mean().item()) + stats["latent_std"] = float(_lat.std().item()) + stats["latent_min"] = float(_lat.min().item()) + stats["latent_max"] = float(_lat.max().item()) + + with open(stats_path, "w", encoding="utf-8") as f: + _json.dump(stats, f, ensure_ascii=False, indent=2) + return f"saved {tag} (raw_mean_rgb={raw_mean})" + + def forward(self, data, inputs=None): + if inputs is None: + inputs = self.forward_preprocess(data) + models = {name: getattr(self.pipe, name) for name in self.pipe.in_iteration_models} + if self.enable_context_memory and "context_latents" in inputs: + return self._training_loss_with_context(**models, **inputs) + inputs = self._ensure_input_latents(inputs, strict=True) + return self.pipe.training_loss(**models, **inputs) + + def _training_loss_with_context(self, **kwargs): + context_latents = kwargs.pop("context_latents", None) + num_context_frames = kwargs.pop("num_context_frames", 0) + models = {k: v for k, v in kwargs.items() if k in self.pipe.in_iteration_models} + inputs = {k: v for k, v in kwargs.items() if k not in self.pipe.in_iteration_models} + if context_latents is not None: + inputs.update({ + "context_latents": context_latents, + "num_context_frames": num_context_frames, + "context_noise_prob": self.context_noise_prob, + "context_noise_std": self.context_noise_std, + "context_attention_weight": getattr(self, "context_attention_weight", 1.0), + "use_anchor_frame": getattr(self, "use_anchor_frame", False), + "context_temporal_decay": getattr(self, "context_temporal_decay", 1.0), + "use_spatial_memory": getattr(self.pipe, "use_spatial_memory", False), + "spatial_memory_tokens": int(getattr(self.pipe, "spatial_memory_tokens", 64) or 64), + "use_spatial_memory_legacy": bool(getattr(self.pipe, "use_spatial_memory_legacy", False)), + "spatial_memory_module": getattr(self.pipe, "spatial_memory_module", None), + "spatial_memory_inject_mode": getattr(self.pipe, "spatial_memory_inject_mode", "concat_text"), + "spatial_memory_readout_module": getattr(self.pipe, "spatial_memory_readout_module", None), + "use_framepack_memory": bool(getattr(self, "use_framepack_memory", False)), + }) + if self.use_moc and self.moc_module is not None: + inputs["use_moc"] = True + inputs["moc_module"] = self.moc_module + if self.context_fixed_noise_std is not None: + inputs["context_fixed_noise_std"] = self.context_fixed_noise_std + inputs = self._ensure_input_latents(inputs, strict=True) + return self.pipe.training_loss(**models, **inputs) + diff --git a/code/src/model_training/transformers_compat.py b/code/src/model_training/transformers_compat.py new file mode 100644 index 0000000000000000000000000000000000000000..aa1d34ba6f83a9c1fca00ffcff4f2edb8aafedd1 --- /dev/null +++ b/code/src/model_training/transformers_compat.py @@ -0,0 +1,25 @@ +def patch_transformers_hybrid_cache() -> None: + """Map missing transformers.HybridCache to DynamicCache for newer peft.""" + try: + import transformers + except Exception: + return + + if hasattr(transformers, "HybridCache"): + return + + lazy_module_cls = type(transformers) + original_getattr = getattr(lazy_module_cls, "__getattr__", None) + if original_getattr is not None and not getattr(lazy_module_cls, "_echo_memory_hybrid_cache_patch", False): + def _compat_getattr(self, name): + if name == "HybridCache": + return original_getattr(self, "DynamicCache") + return original_getattr(self, name) + + lazy_module_cls.__getattr__ = _compat_getattr + lazy_module_cls._echo_memory_hybrid_cache_patch = True + + try: + transformers.HybridCache = transformers.DynamicCache + except Exception: + pass diff --git a/code/tests/test_cgla_wan.py b/code/tests/test_cgla_wan.py new file mode 100644 index 0000000000000000000000000000000000000000..f8a24bb64d1c5a1eba06bfa570c129646355621b --- /dev/null +++ b/code/tests/test_cgla_wan.py @@ -0,0 +1,442 @@ +""" +GPU test: how many Wan 2.1 1.3B weights does the CGLA DiT block load? + +Builds a Wan-2.1-T2V-1.3B-shaped DiT whose every ``DiTBlock`` is replaced by +the CGLA block (``diffsynth/models/memory/u_vit_cgla_blocks.py:: +CGLATransformerBlock`` — the real DFOT SSE-GLA block, sequential attn+mlp, +carrying Wan-shaped submodules ``cross_attn``/``norm1/2/3``/``ffn``/ +``modulation``/``gate``). Loads the Wan 2.1 1.3B checkpoint and reports: + + * which block submodule prefixes loaded successfully (Wan key -> CGLA param) + * which were NOT loaded, split into: + - "unexpected" (Wan has them, CGLA block does not — e.g. Wan's softmax + ``self_attn``), and + - "missing" (CGLA block has them, Wan does not — the new SSE-GLA + attention ``spatial_*``/``temporal_attn``/``noise_write_gate``/``mlp_*``) + * non-block Wan keys (patch_embedding / text_embedding / time_embedding / + final_layer / ...) reported separately (not part of the DiTBlock swap) + * coverage numbers: #keys, #params, % of Wan block params loaded, % of CGLA + block params that are Wan-initialised. + +Wan 2.1 T2V 1.3B config (``diffsynth/models/wan_video_dit.py`` hash +``9269f8db...``): dim=1536, num_heads=12, num_layers=30, ffn_dim=8960, +has_image_input=False, eps=1e-6, patch_size=(1,2,2). These are auto-detected +from the checkpoint where possible (dim/ffn_dim/num_layers/has_image_input); +``num_heads`` is not weight-inferable and uses the known config (12). + +NOTE: this is a weight-LOADING test (no forward pass). Run on a GPU node that +has the Wan base model + the echo-memory env: + + PYTHONPATH=. python3 tests/test_cgla_wan.py + # optional: --ckpt /path/to/diffusion_pytorch_model.safetensors + # optional: --mechanism prope|ucpe (loading result is identical across + # mechanisms; only new params differ) +""" + +from __future__ import annotations + +import os +import sys +from collections import Counter, defaultdict + +_REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if _REPO not in sys.path: + sys.path.insert(0, _REPO) + +import torch +import torch.nn as nn +from safetensors.torch import load_file as safe_load_file + +from diffsynth.models.memory.u_vit_cgla_blocks import CGLATransformerBlock, remap_wan_to_cgla + +# ── Wan 2.1 T2V 1.3B config (diffsynth/models/wan_video_dit.py, hash 9269f8db…) ── +WAN_T2V_1_3B = dict( + dim=1536, + num_heads=12, # not weight-inferable; from the Wan config + num_layers=30, + ffn_dim=8960, + has_image_input=False, + eps=1e-6, + patch_size=(1, 2, 2), +) + +# CGLA block forward-only shape params (do not affect weight shapes/loading): +# 640x352 frame -> VAE /8 -> 80x44 -> DiT patchify /2 -> 40x22 = 880 patches; +# 81 frames -> VAE /4 (+1) -> 21 latent frames. +CGLA_NUM_PATCHES = 880 +CGLA_TEMPORAL_LENGTH = 21 +CGLA_EMB_DIM = 1024 # dfot NormalizeWithCond FiLM emb_dim (CGLA-new param) +CGLA_POSE_DIM = 12 # per-frame RT camera pose +CGLA_HEAD_DIM = 128 # = dim/num_heads = 1536/12 -> SSEGLA key_dim == dim + +DEFAULT_CKPT = "/apdcephfs_zwfy/share_303204533/jiakuihu/checkpoints/Wan2.1-T2V-1.3B/diffusion_pytorch_model.safetensors" + + +def detect_config(sd: dict) -> dict: + """Auto-detect dim / ffn_dim / num_layers / has_image_input from the ckpt.""" + cfg = dict(WAN_T2V_1_3B) + # num_layers: distinct blocks.. prefixes. + blk_ids = sorted({int(k.split(".")[1]) for k in sd if k.startswith("blocks.")}) + if blk_ids: + cfg["num_layers"] = len(blk_ids) + # dim / ffn_dim / has_image_input from block 0 weight shapes. + def _shape(name): + return tuple(sd[name].shape) if name in sd else None + q = _shape("blocks.0.self_attn.q.weight") or _shape("blocks.0.cross_attn.q.weight") + if q is not None: + cfg["dim"] = q[0] + ffn0 = _shape("blocks.0.ffn.0.weight") + if ffn0 is not None: + cfg["ffn_dim"] = ffn0[0] + cfg["has_image_input"] = "blocks.0.cross_attn.k_img.weight" in sd + return cfg + + +def build_cgla_dit(cfg: dict, mechanism: str) -> nn.Module: + """A Wan-shaped DiT whose blocks are CGLATransformerBlock (no VAE/text). + + CGLATransformerBlock is a Wan DiTBlock with CGLA (SSE-GLA) replacing the + self-attention (linear attention on the full flattened token sequence), then + cross-attention to text, then FFN. Wan submodules (cross_attn/norm1/2/3/ + ffn/modulation/gate) + the SSE-GLA q/k/v/o projections load from the Wan + checkpoint; the rest of the SSE-GLA params are new. + """ + use_pose_rope = mechanism in ("prope", "ucpe") + # head_dim = dim // num_heads so SSE-GLA key_dim == dim (Wan's q/k/v/o are + # Linear(dim, dim), loadable into the linear-attention projections). For + # real Wan 1.3B this is 1536 // 12 = 128 (== CGLA_HEAD_DIM). + head_dim = cfg["dim"] // cfg["num_heads"] + + class CGLADiT(nn.Module): + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([ + CGLATransformerBlock( + has_image_input=cfg["has_image_input"], + dim=cfg["dim"], + num_heads=cfg["num_heads"], + ffn_dim=cfg["ffn_dim"], + eps=cfg["eps"], + head_dim=head_dim, + num_sparse_partition=4, + num_writer=1, + num_reader=1, + pose_dim=CGLA_POSE_DIM, + pose_bottleneck=64, + gate_logit_normalizer=16, + gate_low_rank_dim=16, + use_pose_rope=use_pose_rope, + use_pose_gate_mod=False, + layer_idx=i, + emb_dim=CGLA_EMB_DIM, + ) + for i in range(cfg["num_layers"]) + ]) + + return CGLADiT() + + +def _top(prefix_key: str) -> str: + """Block submodule prefix: 'blocks.N....' -> ''; + non-block key -> its first component ('patch_embedding.weight' -> 'patch_embedding').""" + if prefix_key.startswith("blocks."): + parts = prefix_key.split(".") # ["blocks", "", "", ...] + return parts[2] if len(parts) > 2 else parts[-1] + return prefix_key.split(".")[0] + + +def _numel(sd, keys): + return int(sum(sd[k].numel() for k in keys if k in sd)) + + +def _prefix_summary(keys): + """Group keys by their first component; return {prefix: (n_keys, n_params)}.""" + # keys here are full ckpt keys (blocks.N.) — group by [0]. + g = defaultdict(list) + for k in keys: + rest = k.split(".", 2)[2] if k.startswith("blocks.") else k + g[rest.split(".")[0]].append(k) + return g + + +def report(sd, model, cfg, mechanism): + n_blk = cfg["num_layers"] + + # Remap Wan self-attn keys (self_attn.{q,k,v,o}) to CGLA SSE-GLA names + # (self_attn.{q,k,v,o}_proj) so Wan's softmax-attn weights initialise the + # linear-attention projections. After remap, "loaded" reflects this. + sd = remap_wan_to_cgla(sd) + + model_sd = model.state_dict() + ckpt_keys = set(sd.keys()) + model_keys = set(model_sd.keys()) + + # Block-only keys (the DiTBlock swap domain). + ckpt_blk = {k for k in ckpt_keys if k.startswith("blocks.")} + model_blk = {k for k in model_keys if k.startswith("blocks.")} + # Non-block Wan keys (patch_embedding / text_embedding / time_embedding / ...). + ckpt_nonblk = ckpt_keys - ckpt_blk + + missing, unexpected = model.load_state_dict(sd, strict=False) + # Restrict to block keys for the swap analysis (non-block keys are reported + # separately as "not part of the DiTBlock replacement"). + missing_blk = [k for k in missing if k.startswith("blocks.")] + unexpected_blk = [k for k in unexpected if k.startswith("blocks.")] + unexpected_nonblk = [k for k in unexpected if not k.startswith("blocks.") and k in ckpt_nonblk] + + loaded_blk = sorted(model_blk - set(missing_blk)) # in both -> loaded + + # ── param tallies ────────────────────────────────────────────────────── + wan_blk_params = _numel(sd, ckpt_blk) + loaded_params = _numel(sd, loaded_blk) + cgla_blk_params = _numel(model_sd, model_blk) + wan_init_params = loaded_params + + # ── per-prefix grouping (over one block, × num_layers) ───────────────── + def _per_prefix(keys, src): + g = defaultdict(lambda: [0, 0]) # prefix -> [n_keys, n_params] + for k in keys: + pref = _top(k) + g[pref][0] += 1 + g[pref][1] += int(src[k].numel()) if k in src else 0 + return g + + loaded_g = _per_prefix(loaded_blk, sd) + miss_g = _per_prefix(missing_blk, model_sd) + unexp_g = _per_prefix(unexpected_blk, sd) + nonblk_g = _per_prefix(ckpt_nonblk, sd) + + # ── print ────────────────────────────────────────────────────────────── + print("=" * 78) + print("WAN 2.1 1.3B -> CGLA DiT weight-loading report") + print("=" * 78) + print(f"checkpoint : {CKPT}") + print(f"detected Wan config : {cfg}") + print(f"CGLA mechanism : {mechanism}") + print(f" (use_pose_rope={mechanism in ('prope','ucpe')}; loading coverage is") + print(f" identical across cgla/prope/ucpe — only *new* (non-Wan) params differ)") + print(f"num blocks : {n_blk}") + print(f"head_dim (SSEGLA) : {CGLA_HEAD_DIM} (= dim/num_heads -> SSEGLA key_dim == dim)") + print() + + print("-" * 78) + print("1) SUCCESSFULLY LOADED (Wan block key -> CGLA param, name+shape match)") + print("-" * 78) + if loaded_g: + print(f" {'prefix':<22}{'keys (×N blk)':>16}{'params':>16}") + for pref in sorted(loaded_g): + n, p = loaded_g[pref] + print(f" {pref:<22}{n:>16}{p:>16,}") + print(f" {'(total loaded)':<22}{'':>16}{loaded_params:>16,}") + print() + + print("-" * 78) + print("2a) NOT LOADED — unexpected (in Wan ckpt, NOT in CGLA block)") + print("-" * 78) + if unexp_g: + print(f" {'prefix':<22}{'keys (×N blk)':>16}{'params':>16}") + for pref in sorted(unexp_g): + n, p = unexp_g[pref] + print(f" {pref:<22}{n:>16}{p:>16,}") + print(f" {'(total unexpected)':<22}{'':>16}{_numel(sd, unexpected_blk):>16,}") + else: + print(" (none)") + print() + + print("-" * 78) + print("2b) NOT LOADED — missing (CGLA block has, NOT in Wan ckpt = new CGLA params)") + print("-" * 78) + if miss_g: + print(f" {'prefix':<22}{'keys (×N blk)':>16}{'params':>16}") + for pref in sorted(miss_g): + n, p = miss_g[pref] + print(f" {pref:<22}{n:>16}{p:>16,}") + print(f" {'(total missing/new)':<22}{'':>16}{_numel(model_sd, missing_blk):>16,}") + else: + print(" (none)") + print() + + print("-" * 78) + print("3) NON-BLOCK Wan keys (not part of the DiTBlock swap; e.g. patch/text") + print(" /time/final embeddings — the CGLA model has no such submodules)") + print("-" * 78) + if nonblk_g: + print(f" {'prefix':<22}{'keys':>10}{'params':>16}") + for pref in sorted(nonblk_g): + n, p = nonblk_g[pref] + print(f" {pref:<22}{n:>10}{p:>16,}") + else: + print(" (none)") + print() + + # ── coverage summary ─────────────────────────────────────────────────── + wan_total = sum(v.numel() for v in sd.values()) + pct_wan_loaded = 100.0 * loaded_params / wan_blk_params if wan_blk_params else 0.0 + pct_cgla_waninit = 100.0 * wan_init_params / cgla_blk_params if cgla_blk_params else 0.0 + print("=" * 78) + print("COVERAGE SUMMARY") + print("=" * 78) + print(f" Wan ckpt total params : {wan_total:>14,}") + print(f" Wan BLOCK params (blocks.*.* ) : {wan_blk_params:>14,}") + print(f" -> loaded into CGLA block : {loaded_params:>14,} ({pct_wan_loaded:5.2f}% of Wan block params)") + print(f" CGLA block params (new + Wan-shape): {cgla_blk_params:>14,}") + print(f" -> Wan-initialised (loaded) : {wan_init_params:>14,} ({pct_cgla_waninit:5.2f}% of CGLA block params)") + print(f" -> CGLA-new (missing, needs training): {cgla_blk_params - wan_init_params:>14,} " + f"({100.0*(cgla_blk_params-wan_init_params)/cgla_blk_params:5.2f}% of CGLA block params)") + print("=" * 78) + + _report_ssegla_param_audit(model, sd, mechanism) + + +def _self_attn_subparam(key: str): + """For 'blocks.N.self_attn..[...]' -> ''; else None.""" + parts = key.split(".") + if len(parts) >= 3 and parts[0] == "blocks" and parts[2] == "self_attn": + return parts[3] if len(parts) > 3 else None + return None + + +def _report_ssegla_param_audit(model, sd, mechanism): + """Per-param audit of the SSE-GLA ``self_attn`` submodules. + + Answers: for each SSEGLA param (q_proj / k_proj / v_proj / o_proj / + lora_q_proj / lora_k_proj / gk_proj / e_proj / g_proj / o_norm) — LOADED from + Wan (name+shape match via the remap) or REINIT (no Wan key)? And for the + PRoPE/UCPE pose modules (pose_encoder / pose_q_proj / pose_k_proj / + pose_gk_proj / pose_e_proj / pose_rope / pose_gate_mod) — all REINIT, with + their dfot zero-init policy noted (the 'up' projections are zero-init; the + 'down' pose_encoder is default-init). + """ + model_sd = model.state_dict() + # Group by self_attn subparam over all blocks. + def _group(keys, src): + g = defaultdict(lambda: [0, 0]) # subparam -> [n_keys, n_params] + for k in keys: + sp = _self_attn_subparam(k) + if sp is None: + continue + g[sp][0] += 1 + g[sp][1] += int(src[k].numel()) if k in src else 0 + return g + + model_sub = _group(model_sd.keys(), model_sd) # CGLA block params + ckpt_sub = _group(sd.keys(), sd) # Wan (remapped) params + + # Wan self_attn projection names that the remap targets. + WAN_REMAP_TARGETS = {"q_proj", "k_proj", "v_proj", "o_proj"} + # dfot pose 'up'-projection zero-init (see fla/layers/sse.py:352-355, 359-360; + # PoseRoPE net[-1] zero-init at sse.py:98-99). + POSE_ZERO_UP = { + "pose_q_proj", "pose_k_proj", "pose_gk_proj", "pose_e_proj", + "pose_gate_mod", "pose_rope", + } + + print() + print("=" * 78) + print("4) SSE-GLA self_attn PARAM-BY-PARAM (loaded vs reinitialised)") + print("=" * 78) + print(f" {'SSEGLA param':<16}{'Wan source':<20}{'keys (×N blk)':>14}{'params':>14} status") + print(" " + "-" * 74) + total_loaded_sa = 0 + total_reinit_sa = 0 + for sp in sorted(model_sub): + n, p = model_sub[sp] + wan_src = "—" + status = "REINIT (no Wan key)" + if sp in WAN_REMAP_TARGETS: + wan_name = sp.replace("_proj", "") # q_proj -> q + wan_src = f"self_attn.{wan_name}" + if sp in ckpt_sub: + status = "LOADED" + total_loaded_sa += p + else: + status = "REINIT (Wan key absent)" + total_reinit_sa += p + else: + total_reinit_sa += p + note = "" + if sp in POSE_ZERO_UP and mechanism in ("prope", "ucpe"): + note = " [dfot zero-init 'up' proj]" + elif sp == "pose_encoder" and mechanism in ("prope", "ucpe"): + note = " [dfot default-init 'down' proj]" + print(f" {sp:<16}{wan_src:<20}{n:>14}{p:>14,} {status}{note}") + print(" " + "-" * 74) + print(f" self_attn loaded (q/k/v/o_proj from Wan): {total_loaded_sa:>14,}") + print(f" self_attn reinit (GLA gates/LoRA/norm): {total_reinit_sa:>14,}") + print() + + # Pose modules exist only for PRoPE / UCPE. + print("-" * 78) + if mechanism in ("prope", "ucpe"): + print(f"PRoPE/UCPE pose modules (use_pose_rope=True) — ALL reinitialised:") + print(f" {'pose module':<16}{'init policy (dfot)':<40}{'params':>14}") + print(" " + "-" * 70) + pose_rows = [ + ("pose_encoder", "default-init (down proj; NOT zeroed)"), + ("pose_q_proj", "zero-init 'up' proj => pose contributes 0 at step 0"), + ("pose_k_proj", "zero-init 'up' proj"), + ("pose_gk_proj", "zero-init 'up' proj"), + ("pose_e_proj", "zero-init (direct)"), + ("pose_rope", "PoseRoPE net[-1] zero-init (angle MLP)"), + ] + if any(_self_attn_subparam(k) == "pose_gate_mod" for k in model_sd): + pose_rows.append(("pose_gate_mod", "zero-init (weight+bias)")) + for name, policy in pose_rows: + p = model_sub.get(name, [0, 0])[1] + present = "present" if name in model_sub else "ABSENT" + print(f" {name:<16}{policy:<40}{p:>14,} {present}") + print(" => at step 0, pose contributes exactly 0 to q2/k2/gk2/eta (LoRA-style),") + print(" so PRoPE/UCPE is numerically identical to vanilla CGLA until trained.") + else: + print(f"mechanism={mechanism!r}: no pose modules (use_pose_rope=False).") + print("=" * 78) + + +def parse_args(): + import argparse + p = argparse.ArgumentParser(description="CGLA <-> Wan 2.1 1.3B weight-loading test") + p.add_argument("--ckpt", default=DEFAULT_CKPT, + help="path to diffusion_pytorch_model.safetensors") + p.add_argument("--mechanism", default="prope", choices=["cgla", "prope", "ucpe"], + help="CGLA variant (loading coverage is identical across them)") + return p.parse_args() + + +CKPT = DEFAULT_CKPT # set in main() from args + + +def main(): + global CKPT + args = parse_args() + CKPT = args.ckpt + if not os.path.isfile(CKPT): + print(f"ERROR: checkpoint not found: {CKPT}", file=sys.stderr) + print(" set --ckpt /path/to/Wan2.1-T2V-1.3B/diffusion_pytorch_model.safetensors", + file=sys.stderr) + sys.exit(1) + + device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"[cgla_wan] loading checkpoint (CPU): {CKPT}") + sd = safe_load_file(CKPT) + cfg = detect_config(sd) + print(f"[cgla_wan] detected config: {cfg}") + + # Build the CGLA DiT (CPU — loading does not need GPU) and load Wan weights. + model = build_cgla_dit(cfg, args.mechanism) + report(sd, model, cfg, args.mechanism) + + # Confirm the CGLA block constructs on GPU (fla/triton compile at forward, + # not at construction; this just moves params to cuda). + if device == "cuda": + try: + model = model.to(device=device, dtype=torch.bfloat16) + n = sum(p.numel() for p in model.parameters()) + print(f"[cgla_wan] CGLA DiT moved to {device} (bf16); {n:,} params construct OK") + except Exception as e: + print(f"[cgla_wan] WARN: GPU move failed: {e}") + else: + print("[cgla_wan] no CUDA — skipping GPU construct check") + + +if __name__ == "__main__": + main() diff --git a/code/tests/test_context_chunk_utils.py b/code/tests/test_context_chunk_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..522ad17ff8f220b864c03825ffa2cb4ae4d85d5e --- /dev/null +++ b/code/tests/test_context_chunk_utils.py @@ -0,0 +1,46 @@ +"""Quick consistency checks for multichunk-aligned context selection (run: PYTHONPATH=. python3 tests/test_context_chunk_utils.py).""" +import os +import sys + +_repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if _repo not in sys.path: + sys.path.insert(0, _repo) + +from src.model_training.multichunk_sample_utils import ( + context_frames_for_next_chunk, + replay_context_global_indices, + replay_context_actions_from_segment_actions, + prev_chunk_tail_global_indices, +) + + +def test_replay_indices_match_frame_order(): + n, K = 81, 5 + frames = list(range(n)) + picked = context_frames_for_next_chunk(frames, K) + idxs = replay_context_global_indices(n, K) + assert [frames[i] for i in idxs] == picked + + +def test_replay_actions_align(): + n, K = 81, 5 + actions = [[float(i)] * 12 for i in range(n)] + out = replay_context_actions_from_segment_actions(actions, n, K) + idxs = replay_context_global_indices(n, K) + assert out is not None + assert len(out) == len(idxs) + for row, i in zip(out, idxs): + assert row[0] == float(i) + + +def test_prev_chunk_tail_indices(): + assert prev_chunk_tail_global_indices(10, 3) == [7, 8, 9] + assert prev_chunk_tail_global_indices(10, 3, nearest_first=True) == [9, 8, 7] + assert prev_chunk_tail_global_indices(2, 5) is None + + +if __name__ == "__main__": + test_replay_indices_match_frame_order() + test_replay_actions_align() + test_prev_chunk_tail_indices() + print("test_context_chunk_utils: ok") diff --git a/code/tests/test_framepack_memory_align.py b/code/tests/test_framepack_memory_align.py new file mode 100644 index 0000000000000000000000000000000000000000..0531258885c2910755a00ac058a9b7a3dcc67d99 --- /dev/null +++ b/code/tests/test_framepack_memory_align.py @@ -0,0 +1,52 @@ +"""FramePack-Length latent/RT alignment (e.g. K=5, r=4). Run: PYTHONPATH=. python3 tests/test_framepack_memory_align.py""" +import torch + +from diffsynth.models.memory.framepack_length import ( + framepack_align_context_actions_to_latents, + framepack_length_compress_context_latents, +) +from diffsynth.models.memory.framepack_weight import apply_framepack_token_weights + + +def test_k5_r4_latent_and_actions(): + B, C, H, W = 1, 16, 8, 8 + K = 5 + r = 4 + lat = torch.randn(B, C, K, H, W) + out, new_k, K_pad, K_orig = framepack_length_compress_context_latents(lat, r) + assert K_orig == 5 + pad = (r - (K % r)) % r + assert K_pad == K + pad == 8 + assert new_k == 2 + assert out.shape[2] == 2 + ca = torch.randn(K, 12) + aligned = framepack_align_context_actions_to_latents( + ca, K_orig, K_pad, r, device=lat.device, dtype=lat.dtype + ) + assert aligned.shape == (2, 12) + + +def test_framepack_weight_preserves_shape_suffix(): + D = 64 + f, h, w = 5, 2, 2 + num_ctx = 2 + N = f * h * w + x = torch.randn(1, N, D) + y = apply_framepack_token_weights( + x, + num_context_frames=num_ctx, + f=f, + h=h, + w=w, + context_position="suffix", + use_framepack_memory=True, + context_temporal_decay=0.9, + context_attention_weight=1.0, + ) + assert y.shape == x.shape + + +if __name__ == "__main__": + test_k5_r4_latent_and_actions() + test_framepack_weight_preserves_shape_suffix() + print("test_framepack_memory_align: ok") diff --git a/code/tests/test_two_chunk_anchor_readout.py b/code/tests/test_two_chunk_anchor_readout.py new file mode 100644 index 0000000000000000000000000000000000000000..e70a4d5dee07784bfc6a176f872d0b79be2bff8f --- /dev/null +++ b/code/tests/test_two_chunk_anchor_readout.py @@ -0,0 +1,45 @@ +"""Two-chunk action/context alignment + spatial readout sanity tests. + +Run: + PYTHONPATH=. python3 tests/test_two_chunk_anchor_readout.py +""" +import json +import os +import tempfile + +import torch + +from diffsynth.models.memory.spatial_grid_memory import SpatialCrossAttnReadout, apply_spatial_cross_attn_readout +from src.model_training.multichunk_sample_utils import _load_actions_tensor_from_json, _tail_context_actions + + +def test_tail_context_actions_from_left45(): + with tempfile.TemporaryDirectory() as td: + p = os.path.join(td, "action_rotation_left_45.json") + data = { + str(i): [0.0, 0.0, 0.0, 1.0 - 1e-4 * i, -0.01 * i, 0.0, 0.01 * i, 1.0 - 1e-4 * i, 0.0, 0.0, 0.0, 1.0] + for i in range(81) + } + with open(p, "w", encoding="utf-8") as f: + json.dump(data, f) + acts = _load_actions_tensor_from_json(p, device=torch.device("cpu"), dtype=torch.float32) + assert acts is not None and acts.shape == (81, 12) + tail = _tail_context_actions(acts, 5, device=torch.device("cpu"), dtype=torch.float32) + assert tail is not None and tail.shape == (5, 12) + assert torch.allclose(tail[0], acts[-5], atol=1e-6) + assert torch.allclose(tail[-1], acts[-1], atol=1e-6) + + +def test_spatial_cross_attn_readout_shape(): + B, Nt, Nm, D = 1, 128, 64, 64 + x_target = torch.randn(B, Nt, D) + mem = torch.randn(B, Nm, D) + mod = SpatialCrossAttnReadout(dim=D, num_heads=8) + y = apply_spatial_cross_attn_readout(x_target, mem, mod) + assert y.shape == x_target.shape + + +if __name__ == "__main__": + test_tail_context_actions_from_left45() + test_spatial_cross_attn_readout_shape() + print("test_two_chunk_anchor_readout: ok") diff --git a/code/train/README.md b/code/train/README.md new file mode 100644 index 0000000000000000000000000000000000000000..a4886cf1df187df408241f2cf35368f5b81fad4d --- /dev/null +++ b/code/train/README.md @@ -0,0 +1,57 @@ +# Training Recipes + +All launchers assume they are run from the repository root or can derive it from their script path. + +Before training: + +```bash +export WAN_BASE_MODEL=/path/to/Wan2.1-T2V-1.3B +export DATASET_BASE_PATH=/path/to/Context-as-Memory-Dataset +export OUTPUT_BASE_ROOT=$PWD/outputs +``` + +## Memory Baselines + +`train/memory_baselines_basic/` contains the public memory baseline recipes (mechanism mapping: [`../doc/memory_mechanisms.md`](../doc/memory_mechanisms.md)): + +- `run_ablation_no_memory_baseline_two_chunk.sh`: anchor/no-extra-memory reference. +- `run_ablation_framepack_weight_two_chunk.sh`: token weighting without length reduction. +- `run_ablation_framepack_len_r2_two_chunk.sh`: temporal length compression, ratio 2. +- `run_ablation_framepack_len_r4_two_chunk.sh`: temporal length compression, ratio 4. +- `run_ablation_framepack_hybrid_r2_weight_two_chunk.sh`: ratio-2 length compression plus token weighting. +- `run_ablation_framepack_hybrid_r4_weight_two_chunk.sh`: ratio-4 length compression plus token weighting. +- `run_spatial_memory_baseline.sh`: representative spatial memory tokens. +- `run_ablation_spatial_inject_none_two_chunk.sh`: spatial storage with withheld read-out. +- `run_ablation_spatial_concat_text_two_chunk.sh`: spatial memory read through text KV concatenation. +- `run_ablation_spatial_cross_attn_readout_two_chunk.sh`: spatial memory read through dedicated cross-attention. +- `run_videossm_hybrid_baseline.sh`: legacy VideoSSM hybrid memory (`videossm_hybrid.*` checkpoint keys). +- `run_ablation_videossm_hybrid_two_chunk.sh`: legacy VideoSSM with the paper two-chunk monitor. +- `run_ablation_block_wise_ssm_two_chunk.sh`: paper-aligned block-wise SSM recipe (`block_wise_ssm.*` checkpoint keys). +- `run_all_ablations_two_chunk.sh`: convenience launcher for the full ablation set. + +The two-chunk scripts share `common_sampling_two_chunk.sh` and expose `CKPT_INTERVAL`, `TIMESTEP_SHIFT`, `SAMPLING_INTERVAL_STEPS`, `SAMPLING_NUM_INFERENCE_STEPS`, `SAMPLING_HEIGHT`, `SAMPLING_WIDTH`, and `SAMPLING_NUM_FRAMES` as environment overrides. + +## Context Learning + +`train/context_learning/` keeps context-frame recipes: + +- `run_pre_qkv_ctx1.sh` +- `run_pre_qkv_ctx5.sh` +- `run_pre_qkv_ctx20.sh` +- `run_pre_qkv_ctx5_lr_8e5_rt_merge_zero_init_mlp_ssm.sh` +- `run_ctx5_no_action_ablation.sh` + +The shared environment file is `train/_shared/common_env_memory.sh`. It does not set private paths or credentials; configure Weights & Biases in your shell if you use it. + +## Dynamic SpatialVID + +`train/dynamic_spatialvid/` contains the six public dynamic training rows. Set `DATASET_BASE_PATH` to a dynamic `mixed/` root: + +```bash +export DATASET_BASE_PATH=data/dynamic-spatialvid-motion60/mixed +METADATA_NAME=metadata_train.csv bash train/dynamic_spatialvid/run_dyn_spatial_mem.sh +``` + +For local step checks, use `METADATA_NAME=metadata_train_sample_L1.csv MAX_TRAIN_STEPS=1 NUM_WORKERS=0`. + +Dynamic evaluation is TODO; current dynamic support is training and inference. diff --git a/code/train/_shared/common_env_memory.sh b/code/train/_shared/common_env_memory.sh new file mode 100644 index 0000000000000000000000000000000000000000..e4e4a7368b163a5d557c0d1445045f7383b6cc81 --- /dev/null +++ b/code/train/_shared/common_env_memory.sh @@ -0,0 +1,108 @@ +#!/bin/bash +# Shared memory-training environment fragment. +# Source this file from train/*/common_env.sh before setting output_base. +set -euo pipefail + +if [ -f "${CONDA_SH:-}" ]; then + source "${CONDA_SH}" +elif [ -f "$HOME/miniconda3/etc/profile.d/conda.sh" ]; then + source "$HOME/miniconda3/etc/profile.d/conda.sh" +elif [ -f "$HOME/miniconda/etc/profile.d/conda.sh" ]; then + source "$HOME/miniconda/etc/profile.d/conda.sh" +fi +if [ -n "${ECHO_MEMORY_CONDA_ENV:-}" ]; then + conda activate "${ECHO_MEMORY_CONDA_ENV}" 2>/dev/null || true +elif [ -n "${CAM_CONDA_ENV:-}" ]; then + conda activate "${CAM_CONDA_ENV}" 2>/dev/null || true +elif [ -z "${CONDA_DEFAULT_ENV:-}" ]; then + conda activate echo-memory 2>/dev/null || true +fi + +export NCCL_DEBUG="${NCCL_DEBUG:-INFO}" +if [ -z "${NCCL_SOCKET_IFNAME:-}" ]; then + _default_ifname="" + if command -v ip >/dev/null 2>&1; then + _default_ifname="$(ip -o -4 route show to default 2>/dev/null | awk '{print $5; exit}' || true)" + fi + export NCCL_SOCKET_IFNAME="${_default_ifname:-eth0}" +fi +export CONTEXT_POSITION=suffix +export USE_CONCATENATION_INFERENCE=true +export USE_RT_RELATIVE=true + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[1]:-${BASH_SOURCE[0]}}")" && pwd)" +REPO_ROOT="${REPO_ROOT:-$(cd "${SCRIPT_DIR}/../.." && pwd)}" +export PYTHONPATH="${REPO_ROOT}:${PYTHONPATH:-}" + +_first_existing_dir() { + local c + for c in "$@"; do + [ -n "${c}" ] || continue + [ -d "${c}" ] && { echo "${c}"; return 0; } + done + return 1 +} + +_WAN_BASE_DEFAULT="$(_first_existing_dir \ + "${REPO_ROOT}/checkpoints/Wan2.1-T2V-1.3B" \ + "${REPO_ROOT}/checkpoints/wan2.1-t2v-1.3b" \ + "${REPO_ROOT}/models/Wan2.1-T2V-1.3B" \ + || true)" +_WAN_BASE_DEFAULT="${_WAN_BASE_DEFAULT:-}" +WAN_BASE_MODEL="${WAN_BASE_MODEL:-${_WAN_BASE_DEFAULT:-}}" + +_DATASET_BASE_DEFAULT="$(_first_existing_dir \ + "${REPO_ROOT}/data/Context-as-Memory-Dataset" \ + "${REPO_ROOT}/data/Context-as-Memory-Dataset/videos/Context-as-Memory-Dataset" \ + || true)" +_DATASET_BASE_DEFAULT="${_DATASET_BASE_DEFAULT:-}" +DATASET_BASE_PATH="${DATASET_BASE_PATH:-${_DATASET_BASE_DEFAULT:-}}" + +LOG_DIR="${LOG_DIR:-${REPO_ROOT}/logs}" +mkdir -p "${LOG_DIR}" +export WAN_BASE_MODEL DATASET_BASE_PATH LOG_DIR REPO_ROOT + +_require_file() { + local p="$1" + [ -f "${p}" ] || { + echo "[common_env_memory][ERROR] missing file: ${p}" >&2 + echo "[common_env_memory][HINT] set WAN_BASE_MODEL to your Wan2.1 base model directory." >&2 + exit 2 + } +} + +_require_dir() { + local p="$1" + [ -d "${p}" ] || { + echo "[common_env_memory][ERROR] missing dir: ${p}" >&2 + echo "[common_env_memory][HINT] set DATASET_BASE_PATH to your dataset root." >&2 + exit 2 + } +} + +dataset_base_path="${DATASET_BASE_PATH}" +METADATA_NAME="${METADATA_NAME:-metadata_full.csv}" +model_paths="[\"${WAN_BASE_MODEL}/diffusion_pytorch_model.safetensors\",\"${WAN_BASE_MODEL}/models_t5_umt5-xxl-enc-bf16.pth\",\"${WAN_BASE_MODEL}/Wan2.1_VAE.pth\"]" +remove_prefix_in_ckpt="pipe.dit." + +if [ -z "${WAN_BASE_MODEL}" ]; then + echo "[common_env_memory][ERROR] WAN_BASE_MODEL is not set." >&2 + echo "[common_env_memory][HINT] export WAN_BASE_MODEL=/path/to/Wan2.1-T2V-1.3B" >&2 + exit 2 +fi +if [ -z "${DATASET_BASE_PATH}" ]; then + echo "[common_env_memory][ERROR] DATASET_BASE_PATH is not set." >&2 + echo "[common_env_memory][HINT] export DATASET_BASE_PATH=/path/to/Context-as-Memory-Dataset" >&2 + exit 2 +fi +_require_file "${WAN_BASE_MODEL}/diffusion_pytorch_model.safetensors" +_require_file "${WAN_BASE_MODEL}/models_t5_umt5-xxl-enc-bf16.pth" +_require_file "${WAN_BASE_MODEL}/Wan2.1_VAE.pth" +_require_dir "${DATASET_BASE_PATH}" +echo "[common_env_memory] WAN_BASE_MODEL=${WAN_BASE_MODEL}" +echo "[common_env_memory] DATASET_BASE_PATH=${DATASET_BASE_PATH}" + +action_dir="${ACTION_DIR:-${REPO_ROOT}/env}" +sampling_action_path="${action_dir}/action_rotation_left_45.json" +[ ! -f "${sampling_action_path}" ] && (python3 "${action_dir}/generate_rotation_actions.py" 2>/dev/null || true) +cd "${REPO_ROOT}" diff --git a/code/train/context_learning/common_env.sh b/code/train/context_learning/common_env.sh new file mode 100644 index 0000000000000000000000000000000000000000..97de639b388b3ac2264dea53341e70f356973ec0 --- /dev/null +++ b/code/train/context_learning/common_env.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# ctx=5 / ctx=20 + 每帧单独 VAE 注入（无时序降采样），context_actions 每帧一条；baseline 为 run_01_post / run_02_pre_norm +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/../_shared/common_env_memory.sh" +OUTPUT_BASE_ROOT="${OUTPUT_BASE_ROOT:-${REPO_ROOT}/outputs}" +mkdir -p "${OUTPUT_BASE_ROOT}" +output_base="${output_base:-${OUTPUT_BASE_ROOT}/context_learning}" diff --git a/code/train/context_learning/run_ctx5_no_action_ablation.sh b/code/train/context_learning/run_ctx5_no_action_ablation.sh new file mode 100644 index 0000000000000000000000000000000000000000..3e152db7fe15a18a5d355f75e5a951d2bc841b15 --- /dev/null +++ b/code/train/context_learning/run_ctx5_no_action_ablation.sh @@ -0,0 +1,18 @@ +#!/bin/bash +# 消融：无 action 的 ctx=5 训练，用于检查跳变（对比有 action 的 baseline） +# 仅 Context Memory，无 camera_encoder / action 注入 +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --enable_fov_retrieval --fov_top_k 4 --context_memory_frames 5 --context_per_frame_vae --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_ctx_5_no_action_ablation" --trainable_models dit --ckpt_interval 5000 --save_full_model \ + --wandb_run_name "exp1_4_4_ctx_5_no_action_ablation" \ + --enable_context_memory --training_mode context --context_drop_prob 0.1 --cfg_target_only \ + --timestep_shift 5 --enable_video_sampling --sampling_four_prompts --sampling_interval_steps 1000 \ + --verify_high_noise_first_steps 5 --verify_ckpt_step 5 \ + --sampling_num_inference_steps 50 --sampling_negative_prompt "oversaturated colors, overexposed, static, blurry details" \ + --sampling_height 352 --sampling_width 640 --sampling_num_frames 81 --samples_per_epoch 0 \ + 2>&1 | tee "${LOG_DIR}/exp1_4_4_ctx_5_no_action_ablation_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/context_learning/run_pre_qkv_ctx1.sh b/code/train/context_learning/run_pre_qkv_ctx1.sh new file mode 100644 index 0000000000000000000000000000000000000000..f1807d2330f50db0b03eab9ada260840087c28b3 --- /dev/null +++ b/code/train/context_learning/run_pre_qkv_ctx1.sh @@ -0,0 +1,18 @@ +#!/bin/bash +# pre_qkv 系列：action 在第一个 layer norm 之后、3D self-attn 之前与 frame 交互后输入。ctx=1（1 帧 context，无需 per_frame_vae） +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --enable_fov_retrieval --fov_top_k 0 --context_memory_frames 1 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_merged_cam_ctx_1_noise_5_atomic_cam_inject_pre_qkv_seperate_t_r" --trainable_models dit --ckpt_interval 5000 --save_full_model \ + --wandb_run_name "exp1_4_4_ctx_1_noise_5_atomic_cam_inject_pre_qkv" \ + --enable_context_memory --training_mode context --context_drop_prob 0.1 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --timestep_shift 5 --enable_video_sampling --sampling_atomic_left_right --sampling_interval_steps 1000 \ + --verify_high_noise_first_steps 5 --verify_ckpt_step 5 \ + --sampling_num_inference_steps 50 --sampling_negative_prompt "oversaturated colors, overexposed, static, blurry details" \ + --sampling_height 352 --sampling_width 640 --sampling_num_frames 81 --samples_per_epoch 0 --sampling_action_path "${sampling_action_path}" \ + 2>&1 | tee "${LOG_DIR}/exp1_4_4_ctx_1_pre_qkv_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/context_learning/run_pre_qkv_ctx20.sh b/code/train/context_learning/run_pre_qkv_ctx20.sh new file mode 100644 index 0000000000000000000000000000000000000000..0038002899a4b00e8aeba9abafebb94f8ee21c7d --- /dev/null +++ b/code/train/context_learning/run_pre_qkv_ctx20.sh @@ -0,0 +1,18 @@ +#!/bin/bash +# pre_qkv 系列：action 在 norm1 之后、3D self-attn 之前与 frame 交互。ctx=20，每帧单独 VAE、ctx action 按条数输入，target action stride-4 +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --enable_fov_retrieval --fov_top_k 19 --context_memory_frames 20 --context_per_frame_vae --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_merged_cam_ctx_20_noise_5_atomic_cam_inject_pre_qkv_per_frame_vae" --trainable_models dit --ckpt_interval 5000 --save_full_model \ + --wandb_run_name "exp1_4_4_ctx_20_noise_5_atomic_cam_inject_pre_qkv_per_frame_vae" \ + --enable_context_memory --training_mode context --context_drop_prob 0.1 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --timestep_shift 5 --enable_video_sampling --sampling_atomic_left_right --sampling_interval_steps 1000 \ + --verify_high_noise_first_steps 5 --verify_ckpt_step 5 \ + --sampling_num_inference_steps 50 --sampling_negative_prompt "oversaturated colors, overexposed, static, blurry details" \ + --sampling_height 352 --sampling_width 640 --sampling_num_frames 81 --samples_per_epoch 0 --sampling_action_path "${sampling_action_path}" \ + 2>&1 | tee "${LOG_DIR}/exp1_4_4_ctx_20_pre_qkv_per_frame_vae_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/context_learning/run_pre_qkv_ctx5.sh b/code/train/context_learning/run_pre_qkv_ctx5.sh new file mode 100644 index 0000000000000000000000000000000000000000..06dcadcc00d74a0860550d26a4f179c4827b6e89 --- /dev/null +++ b/code/train/context_learning/run_pre_qkv_ctx5.sh @@ -0,0 +1,18 @@ +#!/bin/bash +# pre_qkv 系列：action 在 norm1 之后、3D self-attn 之前与 frame 交互。ctx=5，每帧单独 VAE、ctx action 按条数输入，target action stride-4 +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --enable_fov_retrieval --fov_top_k 4 --context_memory_frames 5 --context_per_frame_vae --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_merged_cam_ctx_5_noise_5_atomic_cam_inject_pre_qkv_per_frame_vae" --trainable_models dit --ckpt_interval 5000 --save_full_model \ + --wandb_run_name "exp1_4_4_ctx_5_noise_5_atomic_cam_inject_pre_qkv_per_frame_vae" \ + --enable_context_memory --training_mode context --context_drop_prob 0.1 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --timestep_shift 5 --enable_video_sampling --sampling_atomic_left_right --sampling_interval_steps 1000 \ + --verify_high_noise_first_steps 5 --verify_ckpt_step 5 \ + --sampling_num_inference_steps 50 --sampling_negative_prompt "oversaturated colors, overexposed, static, blurry details" \ + --sampling_height 352 --sampling_width 640 --sampling_num_frames 81 --samples_per_epoch 0 --sampling_action_path "${sampling_action_path}" \ + 2>&1 | tee "${LOG_DIR}/exp1_4_4_ctx_5_pre_qkv_per_frame_vae_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/context_learning/run_pre_qkv_ctx5_lr_8e5_rt_merge_zero_init_mlp_ssm.sh b/code/train/context_learning/run_pre_qkv_ctx5_lr_8e5_rt_merge_zero_init_mlp_ssm.sh new file mode 100644 index 0000000000000000000000000000000000000000..fafebe883e2f9b5cb59d1cc23df4c187408d13b2 --- /dev/null +++ b/code/train/context_learning/run_pre_qkv_ctx5_lr_8e5_rt_merge_zero_init_mlp_ssm.sh @@ -0,0 +1,20 @@ +#!/bin/bash +# GF-style 整层0初始化 + Diffusion Block SSM：ctx=5 + 合并RT + pre_qkv + 单层MLP + Linear全零初始化 + Block-wise SSM (Long-Context State-Space Video World Models, https://ryanpo.com/ssm_wm/) +# SSM 增强 diffusion block 的 temporal memory，参考 arXiv:2505.20171 +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --enable_fov_retrieval --fov_top_k 4 --context_memory_frames 5 --context_per_frame_vae --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 8e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_merged_cam_ctx_5_pre_qkv_zero_init_mlp_ssm_lr_8e5" --trainable_models dit --ckpt_interval 5000 --save_full_model \ + --wandb_run_name "exp1_4_4_ctx_5_pre_qkv_zero_init_mlp_ssm_lr8e5" \ + --enable_context_memory --training_mode context --context_drop_prob 0.1 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --use_block_wise_ssm --ssm_num_blocks_hint 21 --ssm_every_n_blocks 4 \ + --timestep_shift 5 --enable_video_sampling --sampling_atomic_left_right --sampling_interval_steps 1000 \ + --verify_high_noise_first_steps 5 --verify_ckpt_step 5 \ + --sampling_num_inference_steps 50 --sampling_negative_prompt "oversaturated colors, overexposed, static, blurry details" \ + --sampling_height 352 --sampling_width 640 --sampling_num_frames 81 --samples_per_epoch 0 --sampling_action_path "${sampling_action_path}" \ + 2>&1 | tee "${LOG_DIR}/exp1_4_4_ctx_5_pre_qkv_zero_init_mlp_ssm_lr8e5_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/dynamic_spatialvid/README.md b/code/train/dynamic_spatialvid/README.md new file mode 100644 index 0000000000000000000000000000000000000000..c34adab79f3bf4085645293384cb87e11b619903 --- /dev/null +++ b/code/train/dynamic_spatialvid/README.md @@ -0,0 +1,52 @@ +# Dynamic SpatialVID Training + +Training recipes for the motion-filtered dynamic SpatialVID pool. They mirror the six dynamic rows used for demos and inference: + +| Row | Script | Notes | +| --- | --- | --- | +| Context K=1 | `run_dyn_ctx1.sh` | FOV top-0, 1 context frame | +| Context K=5 | `run_dyn_ctx5.sh` | FOV top-4, per-frame context VAE | +| Context K=20 | `run_dyn_ctx20.sh` | FOV top-19, per-frame context VAE | +| Spatial Memory | `run_dyn_spatial_mem.sh` | 64 spatial memory tokens | +| Block-wise SSM | `run_dyn_block_wise_ssm.sh` | Paper-aligned state-space memory | +| VideoSSM hybrid | `run_dyn_videossm_hybrid.sh` | Legacy temporal-conv baseline | + +Set paths through environment variables: + +```bash +export WAN_BASE_MODEL=/path/to/Wan2.1-T2V-1.3B +export DATASET_BASE_PATH=data/dynamic-spatialvid-motion60/mixed +export OUTPUT_BASE_ROOT=$PWD/outputs/dynamic_spatialvid +``` + +The local validation pool used during development is: + +```bash +export DATASET_BASE_PATH=/pfs/weiyang/DynMemBench-V2/camcl_spatialvid_motion60_ready/mixed +``` + +Public scripts should keep the relative form above. The expected dataset root contains: + +```text +mixed/ +├── frames/L{1,2,3}/{clip_id}/0000.png ... 0080.png +├── jsons/L{1,2,3}/{clip_id}.json +├── overlap_labels/L{1,2,3}/{clip_id}/ +├── metadata_train.csv +├── metadata_train_sample.csv +├── metadata_train_sample_L1.csv +├── metadata_eval.csv +└── metadata_eval_2chunk.csv +``` + +For quick local validation, override the metadata and step count: + +```bash +METADATA_NAME=metadata_train_sample_L1.csv \ +MAX_TRAIN_STEPS=1 \ +PROGRESS_TOTAL_STEPS=30000 \ +NUM_WORKERS=0 \ +bash train/dynamic_spatialvid/run_dyn_block_wise_ssm.sh +``` + +Dynamic evaluation is TODO; current public support covers training and inference. diff --git a/code/train/dynamic_spatialvid/common_env.sh b/code/train/dynamic_spatialvid/common_env.sh new file mode 100644 index 0000000000000000000000000000000000000000..9254c8a5bd5c1008f9e2982f5bae96388b897b0a --- /dev/null +++ b/code/train/dynamic_spatialvid/common_env.sh @@ -0,0 +1,40 @@ +#!/bin/bash +# Shared environment for dynamic SpatialVID training recipes. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +if [ -z "${DATASET_BASE_PATH:-}" ]; then + REPO_ROOT_PRE="${REPO_ROOT:-$(cd "${SCRIPT_DIR}/../.." && pwd)}" + for _d in \ + "${REPO_ROOT_PRE}/data/dynamic-spatialvid-motion60/mixed" \ + "${REPO_ROOT_PRE}/data/DynMemBench-SpatialVID-motion60/mixed" \ + "${REPO_ROOT_PRE}/data/camcl_spatialvid_motion60_ready/mixed"; do + if [ -d "${_d}" ]; then + export DATASET_BASE_PATH="${_d}" + break + fi + done +fi + +source "${SCRIPT_DIR}/../_shared/common_env_memory.sh" + +if [ "${DATASET_BASE_PATH:-}" = "${REPO_ROOT}/data/Context-as-Memory-Dataset" ]; then + for _d in \ + "${REPO_ROOT}/data/dynamic-spatialvid-motion60/mixed" \ + "${REPO_ROOT}/data/DynMemBench-SpatialVID-motion60/mixed" \ + "${REPO_ROOT}/data/camcl_spatialvid_motion60_ready/mixed"; do + if [ -d "${_d}" ]; then + DATASET_BASE_PATH="${_d}" + dataset_base_path="${DATASET_BASE_PATH}" + break + fi + done +fi + +METADATA_NAME="${METADATA_NAME:-metadata_train.csv}" +OUTPUT_BASE_ROOT="${OUTPUT_BASE_ROOT:-${REPO_ROOT}/outputs/dynamic_spatialvid}" +mkdir -p "${OUTPUT_BASE_ROOT}" +output_base="${output_base:-${OUTPUT_BASE_ROOT}/dyn_spatialvid}" + +echo "[dynamic_spatialvid] DATASET_BASE_PATH=${DATASET_BASE_PATH}" +echo "[dynamic_spatialvid] METADATA_NAME=${METADATA_NAME}" diff --git a/code/train/dynamic_spatialvid/run_all_dyn_rows.sh b/code/train/dynamic_spatialvid/run_all_dyn_rows.sh new file mode 100644 index 0000000000000000000000000000000000000000..eb5ed2a8fa0c2510f7573eb4e4eaaede07b0d8dd --- /dev/null +++ b/code/train/dynamic_spatialvid/run_all_dyn_rows.sh @@ -0,0 +1,11 @@ +#!/bin/bash +# Launch all public dynamic SpatialVID rows. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +bash "${SCRIPT_DIR}/run_dyn_ctx1.sh" +bash "${SCRIPT_DIR}/run_dyn_ctx5.sh" +bash "${SCRIPT_DIR}/run_dyn_ctx20.sh" +bash "${SCRIPT_DIR}/run_dyn_spatial_mem.sh" +bash "${SCRIPT_DIR}/run_dyn_block_wise_ssm.sh" +bash "${SCRIPT_DIR}/run_dyn_videossm_hybrid.sh" diff --git a/code/train/dynamic_spatialvid/run_dyn_block_wise_ssm.sh b/code/train/dynamic_spatialvid/run_dyn_block_wise_ssm.sh new file mode 100644 index 0000000000000000000000000000000000000000..208bcd0b22bcdc5867abf52173f298866792268f --- /dev/null +++ b/code/train/dynamic_spatialvid/run_dyn_block_wise_ssm.sh @@ -0,0 +1,25 @@ +#!/bin/bash +# Dynamic SpatialVID: paper-aligned Block-wise SSM. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" + +EXTRA_ARGS=() +[ -n "${MAX_TRAIN_STEPS:-}" ] && EXTRA_ARGS+=(--max_train_steps "${MAX_TRAIN_STEPS}") +[ -n "${PROGRESS_TOTAL_STEPS:-}" ] && EXTRA_ARGS+=(--progress_total_steps "${PROGRESS_TOTAL_STEPS}") + +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --context_source replay --prev_chunk_frames 81 --context_memory_frames 5 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers "${NUM_WORKERS:-16}" \ + --model_paths "${model_paths}" --tokenizer_path "${TOKENIZER_PATH:-${WAN_BASE_MODEL}/google/umt5-xxl}" \ + --learning_rate "${LEARNING_RATE:-5e-5}" --num_epochs "${NUM_EPOCHS:-1}" --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_block_wise_ssm" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "${WANDB_RUN_NAME:-dyn_spatialvid_block_wise_ssm}" \ + --enable_context_memory --training_mode context --context_drop_prob 0.1 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --use_moc --moc_temperature 1.0 \ + --use_block_wise_ssm --ssm_every_n_blocks 4 --ssm_num_blocks_hint 21 \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" \ + "${EXTRA_ARGS[@]}" \ + 2>&1 | tee "${LOG_DIR}/dyn_spatialvid_block_wise_ssm_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/dynamic_spatialvid/run_dyn_ctx1.sh b/code/train/dynamic_spatialvid/run_dyn_ctx1.sh new file mode 100644 index 0000000000000000000000000000000000000000..ea8c548a5e2712ca57b90e61c240d056c5a16cdf --- /dev/null +++ b/code/train/dynamic_spatialvid/run_dyn_ctx1.sh @@ -0,0 +1,23 @@ +#!/bin/bash +# Dynamic SpatialVID: Context K=1 baseline. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" + +EXTRA_ARGS=() +[ -n "${MAX_TRAIN_STEPS:-}" ] && EXTRA_ARGS+=(--max_train_steps "${MAX_TRAIN_STEPS}") +[ -n "${PROGRESS_TOTAL_STEPS:-}" ] && EXTRA_ARGS+=(--progress_total_steps "${PROGRESS_TOTAL_STEPS}") + +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --enable_fov_retrieval --fov_top_k 0 --context_memory_frames 1 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers "${NUM_WORKERS:-16}" \ + --model_paths "${model_paths}" --tokenizer_path "${TOKENIZER_PATH:-${WAN_BASE_MODEL}/google/umt5-xxl}" \ + --learning_rate "${LEARNING_RATE:-5e-5}" --num_epochs "${NUM_EPOCHS:-1}" --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_ctx1" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "${WANDB_RUN_NAME:-dyn_spatialvid_ctx1}" \ + --enable_context_memory --training_mode context --context_drop_prob 0.1 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --timestep_shift "${TIMESTEP_SHIFT:-5}" \ + "${EXTRA_ARGS[@]}" \ + 2>&1 | tee "${LOG_DIR}/dyn_spatialvid_ctx1_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/dynamic_spatialvid/run_dyn_ctx20.sh b/code/train/dynamic_spatialvid/run_dyn_ctx20.sh new file mode 100644 index 0000000000000000000000000000000000000000..684997f53a620a3b6301549211587fb9500283c1 --- /dev/null +++ b/code/train/dynamic_spatialvid/run_dyn_ctx20.sh @@ -0,0 +1,23 @@ +#!/bin/bash +# Dynamic SpatialVID: Context K=20 baseline. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" + +EXTRA_ARGS=() +[ -n "${MAX_TRAIN_STEPS:-}" ] && EXTRA_ARGS+=(--max_train_steps "${MAX_TRAIN_STEPS}") +[ -n "${PROGRESS_TOTAL_STEPS:-}" ] && EXTRA_ARGS+=(--progress_total_steps "${PROGRESS_TOTAL_STEPS}") + +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --enable_fov_retrieval --fov_top_k 19 --context_memory_frames 20 --context_per_frame_vae --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers "${NUM_WORKERS:-16}" \ + --model_paths "${model_paths}" --tokenizer_path "${TOKENIZER_PATH:-${WAN_BASE_MODEL}/google/umt5-xxl}" \ + --learning_rate "${LEARNING_RATE:-5e-5}" --num_epochs "${NUM_EPOCHS:-1}" --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_ctx20" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "${WANDB_RUN_NAME:-dyn_spatialvid_ctx20}" \ + --enable_context_memory --training_mode context --context_drop_prob 0.1 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --timestep_shift "${TIMESTEP_SHIFT:-5}" \ + "${EXTRA_ARGS[@]}" \ + 2>&1 | tee "${LOG_DIR}/dyn_spatialvid_ctx20_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/dynamic_spatialvid/run_dyn_ctx5.sh b/code/train/dynamic_spatialvid/run_dyn_ctx5.sh new file mode 100644 index 0000000000000000000000000000000000000000..5f0ab15e7f8063f781c4ee1af8965349c7f21822 --- /dev/null +++ b/code/train/dynamic_spatialvid/run_dyn_ctx5.sh @@ -0,0 +1,23 @@ +#!/bin/bash +# Dynamic SpatialVID: Context K=5 baseline. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" + +EXTRA_ARGS=() +[ -n "${MAX_TRAIN_STEPS:-}" ] && EXTRA_ARGS+=(--max_train_steps "${MAX_TRAIN_STEPS}") +[ -n "${PROGRESS_TOTAL_STEPS:-}" ] && EXTRA_ARGS+=(--progress_total_steps "${PROGRESS_TOTAL_STEPS}") + +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --enable_fov_retrieval --fov_top_k 4 --context_memory_frames 5 --context_per_frame_vae --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers "${NUM_WORKERS:-16}" \ + --model_paths "${model_paths}" --tokenizer_path "${TOKENIZER_PATH:-${WAN_BASE_MODEL}/google/umt5-xxl}" \ + --learning_rate "${LEARNING_RATE:-5e-5}" --num_epochs "${NUM_EPOCHS:-1}" --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_ctx5" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "${WANDB_RUN_NAME:-dyn_spatialvid_ctx5}" \ + --enable_context_memory --training_mode context --context_drop_prob 0.1 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --timestep_shift "${TIMESTEP_SHIFT:-5}" \ + "${EXTRA_ARGS[@]}" \ + 2>&1 | tee "${LOG_DIR}/dyn_spatialvid_ctx5_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/dynamic_spatialvid/run_dyn_spatial_mem.sh b/code/train/dynamic_spatialvid/run_dyn_spatial_mem.sh new file mode 100644 index 0000000000000000000000000000000000000000..01795d1a533a0d8d0607891f25db9dbb5c04bf46 --- /dev/null +++ b/code/train/dynamic_spatialvid/run_dyn_spatial_mem.sh @@ -0,0 +1,25 @@ +#!/bin/bash +# Dynamic SpatialVID: Spatial Memory baseline. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" + +EXTRA_ARGS=() +[ -n "${MAX_TRAIN_STEPS:-}" ] && EXTRA_ARGS+=(--max_train_steps "${MAX_TRAIN_STEPS}") +[ -n "${PROGRESS_TOTAL_STEPS:-}" ] && EXTRA_ARGS+=(--progress_total_steps "${PROGRESS_TOTAL_STEPS}") + +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --context_source prev_chunk_tail --context_memory_frames 1 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers "${NUM_WORKERS:-16}" \ + --model_paths "${model_paths}" --tokenizer_path "${TOKENIZER_PATH:-${WAN_BASE_MODEL}/google/umt5-xxl}" \ + --learning_rate "${LEARNING_RATE:-5e-5}" --num_epochs "${NUM_EPOCHS:-1}" --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_spatial_mem" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "${WANDB_RUN_NAME:-dyn_spatialvid_spatial_mem}" \ + --enable_context_memory --training_mode context --context_drop_prob 0.1 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --use_moc --moc_temperature 1.0 \ + --use_spatial_memory --spatial_memory_tokens 64 \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" \ + "${EXTRA_ARGS[@]}" \ + 2>&1 | tee "${LOG_DIR}/dyn_spatialvid_spatial_mem_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/dynamic_spatialvid/run_dyn_videossm_hybrid.sh b/code/train/dynamic_spatialvid/run_dyn_videossm_hybrid.sh new file mode 100644 index 0000000000000000000000000000000000000000..613bae53acb3f9c9738343b1a00400e358890777 --- /dev/null +++ b/code/train/dynamic_spatialvid/run_dyn_videossm_hybrid.sh @@ -0,0 +1,25 @@ +#!/bin/bash +# Dynamic SpatialVID: legacy VideoSSM hybrid baseline. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" + +EXTRA_ARGS=() +[ -n "${MAX_TRAIN_STEPS:-}" ] && EXTRA_ARGS+=(--max_train_steps "${MAX_TRAIN_STEPS}") +[ -n "${PROGRESS_TOTAL_STEPS:-}" ] && EXTRA_ARGS+=(--progress_total_steps "${PROGRESS_TOTAL_STEPS}") + +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --enable_fov_retrieval --fov_top_k 4 --context_memory_frames 5 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers "${NUM_WORKERS:-16}" \ + --model_paths "${model_paths}" --tokenizer_path "${TOKENIZER_PATH:-${WAN_BASE_MODEL}/google/umt5-xxl}" \ + --learning_rate "${LEARNING_RATE:-5e-5}" --num_epochs "${NUM_EPOCHS:-1}" --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_videossm_hybrid" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "${WANDB_RUN_NAME:-dyn_spatialvid_videossm_hybrid}" \ + --enable_context_memory --training_mode context --context_drop_prob 0.1 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --use_moc --moc_temperature 1.0 \ + --use_videossm_hybrid --videossm_every_n_blocks 4 --videossm_kernel_size 3 --videossm_expand 2 \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" \ + "${EXTRA_ARGS[@]}" \ + 2>&1 | tee "${LOG_DIR}/dyn_spatialvid_videossm_hybrid_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/memory_baselines_basic/README.md b/code/train/memory_baselines_basic/README.md new file mode 100644 index 0000000000000000000000000000000000000000..832d8a0b8d1b2ebfc833ada41e1ddcea609ae991 --- /dev/null +++ b/code/train/memory_baselines_basic/README.md @@ -0,0 +1,101 @@ +# Memory Baselines Basic: FramePack, Spatial, and State-Space + +This folder contains the public training recipes for the paper's controlled memory-design matrix. See `../../doc/memory_mechanisms.md` for the concise paper-row to implementation map. The scripts vary only the memory/context profile while keeping the backbone, optimizer, action conditioning, and evaluation interface aligned. + +## Paper Rows and Code Mapping + +| Paper row | Mechanism | Main implementation | +|-----------|-----------|---------------------| +| FramePack-Weight | Per-frame temporal decay and global scaling over context tokens; context length is unchanged. | `diffsynth/models/memory/framepack_weight.py` | +| FramePack-Length | Temporal mean pooling over context latents with matched RT-action padding and pooling. | `diffsynth/models/memory/framepack_length.py` | +| Hybrid FramePack | Length compression plus token weighting. | `wan_video_new.py` memory path plus FramePack helpers | +| Token-grid baseline (`spatial_mem`) | Time-mean context summary to learned spatial grid tokens. This row has no depth/3D reconstruction. | `diffsynth/models/memory/spatial_grid_memory.py` | +| Geometry-grounded Spatial Memory | TSDF-fused static point-cloud renders are VAE-encoded and summarized into geometry conditioning tokens. | `diffsynth/models/memory/geometry_spatial_memory.py` | +| Block-wise SSM | Paper-aligned recurrent state inside selected DiT blocks. | `diffsynth/models/memory/block_wise_ssm.py` + `--use_block_wise_ssm` | +| VideoSSM hybrid | Legacy lightweight temporal-convolution state-space baseline; kept separate from Block-wise SSM. | `diffsynth/models/memory/videossm_hybrid.py` + `--use_videossm_hybrid` | + +## Two-Chunk Ablation Scripts + +The two-chunk scripts source `common_env.sh` and `common_sampling_two_chunk.sh`. The monitor uses `left_45` followed by `right_45`, writes `sampling_videos/step_*_two_chunk_memory*.mp4`, and stores metadata beside the videos. + +| Script | Purpose | +|--------|---------| +| `run_ablation_no_memory_baseline_two_chunk.sh` | Anchor/no-extra-memory reference. | +| `run_ablation_framepack_weight_two_chunk.sh` | FramePack token weighting. | +| `run_ablation_framepack_len_r2_two_chunk.sh` | Length compression with ratio 2. | +| `run_ablation_framepack_len_r4_two_chunk.sh` | Length compression with ratio 4. | +| `run_ablation_framepack_hybrid_r2_weight_two_chunk.sh` | Ratio-2 length compression plus token weighting. | +| `run_ablation_framepack_hybrid_r4_weight_two_chunk.sh` | Ratio-4 length compression plus token weighting. | +| `run_ablation_spatial_inject_none_two_chunk.sh` | Spatial tokens are stored but not injected. | +| `run_ablation_spatial_concat_text_two_chunk.sh` | Spatial tokens are appended to text cross-attention keys/values. | +| `run_ablation_spatial_cross_attn_readout_two_chunk.sh` | Spatial tokens are read through a dedicated cross-attention read-out. | +| `run_ablation_videossm_hybrid_two_chunk.sh` | Legacy VideoSSM hybrid with the two-chunk monitor. | +| `run_ablation_block_wise_ssm_two_chunk.sh` | Paper-aligned block-wise SSM. | +| `run_all_ablations_two_chunk.sh` | Sequential launcher for the full ablation set. | + +## Representative Baselines + +The non-ablation baseline scripts are still useful for representative rows and quick training checks: + +- `run_framepack_baseline.sh`: FramePack weight-only baseline. +- `run_framepack_lencompress_r2.sh`: FramePack length compression ratio 2. +- `run_framepack_lencompress_r4.sh`: FramePack length compression ratio 4. +- `run_spatial_memory_baseline.sh`: representative Spatial Memory baseline. +- `run_geometry_spatial_memory_baseline.sh`: geometry-grounded adaptation of + [Video World Models with Long-term Spatial Memory](https://arxiv.org/abs/2506.05284). + It requires a metadata `geometry_memory` column containing paths to rendered + static point-cloud videos such as the reference implementation's + `Vid_masktarget.mp4`. +- `run_videossm_hybrid_baseline.sh`: legacy VideoSSM hybrid baseline. + +These scripts expose common overrides through environment variables: `CKPT_INTERVAL`, `TIMESTEP_SHIFT`, `SAMPLING_INTERVAL_STEPS`, `SAMPLING_NUM_INFERENCE_STEPS`, `SAMPLING_HEIGHT`, `SAMPLING_WIDTH`, and `SAMPLING_NUM_FRAMES`. + +## Context Construction + +- `--context_source fov`: FOV-overlap retrieval over historical frames. +- `--context_source replay`: replay-style context construction aligned with `env/run_replay_loop_two_chunk.py`. +- `--context_source prev_chunk_tail`: continuous frames from `[start_frame - N, start_frame)` on disk. + +The shared implementation lives in `src/model_training/context_chunk_utils.py`. + +## Evaluation Alignment + +Evaluation scripts source `env/eval_infer_alignment_env.sh` and call `env/memory_baseline_runtime.py` to infer the correct runtime memory flags from checkpoint paths. Keep output suffixes stable if you add a new script, or update `env/memory_baseline_runtime.py` so evaluation can recover the matching memory profile. + +Generated training monitor videos are useful for fast visual checks, but paper-quality revisit panels should be produced through `eval/v2/revisit_suite`, which stores first frames, revisit-tail frames, change maps, and generated MP4 files for each case. + +## Preparing Geometry Memory + +The geometry extractor is not `SpatialGridMemory`. Follow the official +[`spmem/spmem`](https://github.com/spmem/spmem) preprocessing path: + +1. Recover RGB, metric depth, camera intrinsics, and camera-to-world poses + (Mega-SaM for offline training data; CUT3R is used by the reference work for + recurrent online reconstruction). +2. Run TSDF fusion and render the static point cloud along the target camera + trajectory. The reference script writes `Vid_masktarget.mp4`. + +```bash +SPMEM_ROOT=/path/to/spmem \ +INPUT_NPZ=/path/to/reconstructed_clip.npz \ +CLIP_NAME=sample_id \ +GEOMETRY_OUTPUT_ROOT=/path/to/tsdf/outputs \ +bash scripts/run_spmem_tsdf_preprocess.sh +``` + +3. Add those video paths to metadata: + +```bash +python scripts/add_geometry_memory_column.py \ + --metadata /path/to/metadata_full.csv \ + --geometry_root /path/to/tsdf/outputs \ + --output /path/to/metadata_geometry.csv +``` + +4. Train with: + +```bash +METADATA_NAME=metadata_geometry.csv \ +GEOMETRY_MEMORY_ROOT=/path/to/tsdf/outputs \ +bash train/memory_baselines_basic/run_geometry_spatial_memory_baseline.sh +``` diff --git a/code/train/memory_baselines_basic/common_env.sh b/code/train/memory_baselines_basic/common_env.sh new file mode 100644 index 0000000000000000000000000000000000000000..240ad3453a8d26ed5a7db905cb91543168ddfa41 --- /dev/null +++ b/code/train/memory_baselines_basic/common_env.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# memory_baselines_basic: FAR/FramePack / Spatial Memory / VideoSSM baseline 共用环境 +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/../_shared/common_env_memory.sh" +OUTPUT_BASE_ROOT="${OUTPUT_BASE_ROOT:-${REPO_ROOT}/outputs}" +mkdir -p "${OUTPUT_BASE_ROOT}" +output_base="${output_base:-${OUTPUT_BASE_ROOT}/memory_baselines_basic}" diff --git a/code/train/memory_baselines_basic/common_sampling_two_chunk.sh b/code/train/memory_baselines_basic/common_sampling_two_chunk.sh new file mode 100644 index 0000000000000000000000000000000000000000..6ce6b70151bd4683e2552b83a03ae87929600c95 --- /dev/null +++ b/code/train/memory_baselines_basic/common_sampling_two_chunk.sh @@ -0,0 +1,29 @@ +#!/bin/bash +# Source after common_env.sh: shared train-side video sampling = 2-chunk monitor (eval-aligned). +# Replaces --sampling_atomic_left_right. Mutually exclusive with four_prompts / two_prompts in ModelLogger. +# shellcheck disable=SC2034 +left45_action_path="$(dirname "${sampling_action_path}")/action_rotation_left_45.json" +right45_action_path="$(dirname "${sampling_action_path}")/action_rotation_right_45.json" +SAMPLING_TWO_CHUNK_FLAGS=( + --enable_video_sampling + --sampling_two_chunk_memory + --sampling_interval_steps + "${SAMPLING_INTERVAL_STEPS:-1000}" + --sampling_two_chunk_action_path + "${left45_action_path}" + --sampling_num_inference_steps + "${SAMPLING_NUM_INFERENCE_STEPS:-50}" + --sampling_negative_prompt + "oversaturated colors, overexposed, static, blurry details" + --sampling_height + "${SAMPLING_HEIGHT:-352}" + --sampling_width + "${SAMPLING_WIDTH:-640}" + --sampling_num_frames + "${SAMPLING_NUM_FRAMES:-81}" + --samples_per_epoch + "0" + --sampling_action_path + "${right45_action_path}" + --use_anchor_frame +) diff --git a/code/train/memory_baselines_basic/run_ablation_block_wise_ssm_two_chunk.sh b/code/train/memory_baselines_basic/run_ablation_block_wise_ssm_two_chunk.sh new file mode 100644 index 0000000000000000000000000000000000000000..ceb2600dd1420be6a25b1e07de267311ac117324 --- /dev/null +++ b/code/train/memory_baselines_basic/run_ablation_block_wise_ssm_two_chunk.sh @@ -0,0 +1,20 @@ +#!/bin/bash +# Ablation: paper-style block-wise SSM (arXiv:2505.20171) + 2-chunk monitor. Prefer over legacy VideoSSM hybrid for SSM table. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +source "${SCRIPT_DIR}/common_sampling_two_chunk.sh" +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --context_source replay --prev_chunk_frames 81 --context_memory_frames 5 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_abl_block_wise_ssm_two_chunk" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "abl_block_wise_ssm_two_chunk" \ + --enable_context_memory --training_mode context --context_drop_prob 0.1 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --use_moc --moc_temperature 1.0 \ + --use_block_wise_ssm --ssm_every_n_blocks 4 --ssm_num_blocks_hint 21 \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" \ + "${SAMPLING_TWO_CHUNK_FLAGS[@]}" \ + 2>&1 | tee "${LOG_DIR}/abl_block_wise_ssm_two_chunk_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/memory_baselines_basic/run_ablation_cgla_memory_two_chunk.sh b/code/train/memory_baselines_basic/run_ablation_cgla_memory_two_chunk.sh new file mode 100644 index 0000000000000000000000000000000000000000..960ecea7c66ea3aa2837268321104720ad207ad0 --- /dev/null +++ b/code/train/memory_baselines_basic/run_ablation_cgla_memory_two_chunk.sh @@ -0,0 +1,24 @@ +#!/bin/bash +# Ablation: Camera-Guided Linear Attention (CGLA) memory + 2-chunk monitor. +# Pose-conditioned gated linear attention (DFOT's SSEGLA) as the per-block +# temporal-memory pathway. Mirrors run_ablation_block_wise_ssm_two_chunk.sh; +# only the memory pathway differs (controlled ablation). +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +source "${SCRIPT_DIR}/common_sampling_two_chunk.sh" +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --enable_fov_retrieval --fov_top_k 4 --context_memory_frames 5 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_abl_cgla_memory_two_chunk" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "abl_cgla_memory_two_chunk" \ + --enable_context_memory --training_mode context --condition_t2v_ratio 0.10 --condition_i2v_ratio 0.10 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --spike_threshold "${SPIKE_THRESHOLD:-15.0}" \ + --use_moc --moc_temperature 1.0 \ + --use_cgla_memory --cgla_every_n_blocks 4 --cgla_aux_loss_weight 0.01 --cgla_mechanism cgla \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" \ + "${SAMPLING_TWO_CHUNK_FLAGS[@]}" \ + 2>&1 | tee "${LOG_DIR}/abl_cgla_memory_two_chunk_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/memory_baselines_basic/run_ablation_framepack_hybrid_r2_weight_two_chunk.sh b/code/train/memory_baselines_basic/run_ablation_framepack_hybrid_r2_weight_two_chunk.sh new file mode 100644 index 0000000000000000000000000000000000000000..aaa2f54d8b5d521c9fbe5a7633c3c96ef5d55c62 --- /dev/null +++ b/code/train/memory_baselines_basic/run_ablation_framepack_hybrid_r2_weight_two_chunk.sh @@ -0,0 +1,24 @@ +#!/bin/bash +# Ablation: hybrid FramePack (length r=2 + weight decay) + 2-chunk monitor. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +source "${SCRIPT_DIR}/common_sampling_two_chunk.sh" +export CONTEXT_POSITION=suffix +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --context_source prev_chunk_tail --context_memory_frames 81 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_abl_framepack_hybrid_r2_weight_two_chunk" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "abl_framepack_hybrid_r2_weight_two_chunk" \ + --enable_context_memory --training_mode context --condition_t2v_ratio 0.10 --condition_i2v_ratio 0.10 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --spike_threshold "${SPIKE_THRESHOLD:-15.0}" \ + --use_moc --moc_temperature 1.0 \ + --use_framepack_memory --context_temporal_decay 0.95 --context_attention_weight 1.1 \ + --use_framepack_length_compress --framepack_ratio 2 \ + --framepack_length_strategy packed_multiscale --framepack_multiscale_w2 0.25 --framepack_multiscale_w4 0.15 \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" \ + "${SAMPLING_TWO_CHUNK_FLAGS[@]}" \ + 2>&1 | tee "${LOG_DIR}/abl_framepack_hybrid_r2_weight_two_chunk_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/memory_baselines_basic/run_ablation_framepack_hybrid_r4_weight_two_chunk.sh b/code/train/memory_baselines_basic/run_ablation_framepack_hybrid_r4_weight_two_chunk.sh new file mode 100644 index 0000000000000000000000000000000000000000..423cec8e03054cf4073dd350b52209510943189b --- /dev/null +++ b/code/train/memory_baselines_basic/run_ablation_framepack_hybrid_r4_weight_two_chunk.sh @@ -0,0 +1,24 @@ +#!/bin/bash +# Ablation: hybrid FramePack (length r=4 + weight decay) + 2-chunk monitor. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +source "${SCRIPT_DIR}/common_sampling_two_chunk.sh" +export CONTEXT_POSITION=suffix +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --context_source prev_chunk_tail --context_memory_frames 81 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_abl_framepack_hybrid_r4_weight_two_chunk" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "abl_framepack_hybrid_r4_weight_two_chunk" \ + --enable_context_memory --training_mode context --condition_t2v_ratio 0.10 --condition_i2v_ratio 0.10 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --spike_threshold "${SPIKE_THRESHOLD:-15.0}" \ + --use_moc --moc_temperature 1.0 \ + --use_framepack_memory --context_temporal_decay 0.95 --context_attention_weight 1.1 \ + --use_framepack_length_compress --framepack_ratio 4 \ + --framepack_length_strategy packed_multiscale --framepack_multiscale_w2 0.25 --framepack_multiscale_w4 0.15 \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" \ + "${SAMPLING_TWO_CHUNK_FLAGS[@]}" \ + 2>&1 | tee "${LOG_DIR}/abl_framepack_hybrid_r4_weight_two_chunk_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/memory_baselines_basic/run_ablation_framepack_len_r2_two_chunk.sh b/code/train/memory_baselines_basic/run_ablation_framepack_len_r2_two_chunk.sh new file mode 100644 index 0000000000000000000000000000000000000000..cd43a96f1771b383b8a055c5b32fdd79deecf42e --- /dev/null +++ b/code/train/memory_baselines_basic/run_ablation_framepack_len_r2_two_chunk.sh @@ -0,0 +1,23 @@ +#!/bin/bash +# Ablation: FramePack-Length r=2 + 2-chunk monitor. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +source "${SCRIPT_DIR}/common_sampling_two_chunk.sh" +export CONTEXT_POSITION=suffix +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --context_source prev_chunk_tail --context_memory_frames 81 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_abl_framepack_len_r2_two_chunk" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "abl_framepack_len_r2_two_chunk" \ + --enable_context_memory --training_mode context --condition_t2v_ratio 0.10 --condition_i2v_ratio 0.10 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --spike_threshold "${SPIKE_THRESHOLD:-15.0}" \ + --use_moc --moc_temperature 1.0 \ + --use_framepack_length_compress --framepack_ratio 2 \ + --framepack_length_strategy packed_multiscale --framepack_multiscale_w2 0.25 --framepack_multiscale_w4 0.15 \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" \ + "${SAMPLING_TWO_CHUNK_FLAGS[@]}" \ + 2>&1 | tee "${LOG_DIR}/abl_framepack_len_r2_two_chunk_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/memory_baselines_basic/run_ablation_framepack_len_r4_two_chunk.sh b/code/train/memory_baselines_basic/run_ablation_framepack_len_r4_two_chunk.sh new file mode 100644 index 0000000000000000000000000000000000000000..f886e1ea5613eda389495488b36ceb1b7be6bb73 --- /dev/null +++ b/code/train/memory_baselines_basic/run_ablation_framepack_len_r4_two_chunk.sh @@ -0,0 +1,23 @@ +#!/bin/bash +# Ablation: FramePack-Length r=4 + 2-chunk monitor. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +source "${SCRIPT_DIR}/common_sampling_two_chunk.sh" +export CONTEXT_POSITION=suffix +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --context_source prev_chunk_tail --context_memory_frames 81 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_abl_framepack_len_r4_two_chunk" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "abl_framepack_len_r4_two_chunk" \ + --enable_context_memory --training_mode context --condition_t2v_ratio 0.10 --condition_i2v_ratio 0.10 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --spike_threshold "${SPIKE_THRESHOLD:-15.0}" \ + --use_moc --moc_temperature 1.0 \ + --use_framepack_length_compress --framepack_ratio 4 \ + --framepack_length_strategy packed_multiscale --framepack_multiscale_w2 0.25 --framepack_multiscale_w4 0.15 \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" \ + "${SAMPLING_TWO_CHUNK_FLAGS[@]}" \ + 2>&1 | tee "${LOG_DIR}/abl_framepack_len_r4_two_chunk_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/memory_baselines_basic/run_ablation_framepack_weight_two_chunk.sh b/code/train/memory_baselines_basic/run_ablation_framepack_weight_two_chunk.sh new file mode 100644 index 0000000000000000000000000000000000000000..85b47df6c1a53b233bc917771d0b1ef5e3984f26 --- /dev/null +++ b/code/train/memory_baselines_basic/run_ablation_framepack_weight_two_chunk.sh @@ -0,0 +1,22 @@ +#!/bin/bash +# Ablation: FramePack-Weight + 2-chunk train monitor. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +source "${SCRIPT_DIR}/common_sampling_two_chunk.sh" +export CONTEXT_POSITION=suffix +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --context_source prev_chunk_tail --context_memory_frames 81 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_abl_framepack_weight_two_chunk" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "abl_framepack_weight_two_chunk" \ + --enable_context_memory --training_mode context --condition_t2v_ratio 0.10 --condition_i2v_ratio 0.10 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --spike_threshold "${SPIKE_THRESHOLD:-15.0}" \ + --use_moc --moc_temperature 1.0 \ + --use_framepack_memory --context_temporal_decay 0.9 --context_attention_weight 1.0 \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" \ + "${SAMPLING_TWO_CHUNK_FLAGS[@]}" \ + 2>&1 | tee "${LOG_DIR}/abl_framepack_weight_two_chunk_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/memory_baselines_basic/run_ablation_no_memory_baseline_two_chunk.sh b/code/train/memory_baselines_basic/run_ablation_no_memory_baseline_two_chunk.sh new file mode 100644 index 0000000000000000000000000000000000000000..d66537a6fe05aa2b4737d1a39d4dc617f8b0b175 --- /dev/null +++ b/code/train/memory_baselines_basic/run_ablation_no_memory_baseline_two_chunk.sh @@ -0,0 +1,19 @@ +#!/bin/bash +# Ablation control: no FramePack / no spatial / no SSM extras; replay context + 2-chunk monitor only. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +source "${SCRIPT_DIR}/common_sampling_two_chunk.sh" +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --context_source replay --prev_chunk_frames 81 --context_memory_frames 5 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_abl_no_memory_extra_two_chunk" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "abl_no_memory_extra_two_chunk" \ + --enable_context_memory --training_mode context --context_drop_prob 0.1 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --use_moc --moc_temperature 1.0 \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" \ + "${SAMPLING_TWO_CHUNK_FLAGS[@]}" \ + 2>&1 | tee "${LOG_DIR}/abl_no_memory_extra_two_chunk_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/memory_baselines_basic/run_ablation_prope_memory_two_chunk.sh b/code/train/memory_baselines_basic/run_ablation_prope_memory_two_chunk.sh new file mode 100644 index 0000000000000000000000000000000000000000..6a71a5ad9202bbfa1158662eae565bc75cbfb878 --- /dev/null +++ b/code/train/memory_baselines_basic/run_ablation_prope_memory_two_chunk.sh @@ -0,0 +1,24 @@ +#!/bin/bash +# Ablation: PRoPE — CGLA (SSE-GLA) with camera-pose rotary position encoding +# (dfot `use_pose_rope` / PoseRoPE): q/k rotated by learned per-token camera-pose +# angles, so the GLA bilinear form depends on the RELATIVE camera pose. +# Same SSE-GLA backbone + contract + hyperparameters as run_ablation_cgla_memory_two_chunk.sh; +# only --cgla_mechanism prope differs (controlled ablation). +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +source "${SCRIPT_DIR}/common_sampling_two_chunk.sh" +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --context_source replay --prev_chunk_frames 81 --context_memory_frames 5 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_abl_prope_memory_two_chunk" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "abl_prope_memory_two_chunk" \ + --enable_context_memory --training_mode context --context_drop_prob 0.1 --cfg_target_only \ + --train_cam_pose --add_action_attn \ + --use_moc --moc_temperature 1.0 \ + --use_cgla_memory --cgla_every_n_blocks 1 --cgla_aux_loss_weight 0.01 --cgla_mechanism prope \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" \ + "${SAMPLING_TWO_CHUNK_FLAGS[@]}" \ + 2>&1 | tee "${LOG_DIR}/abl_prope_memory_two_chunk_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/memory_baselines_basic/run_ablation_prope_memory_two_chunk.sh.p0bak b/code/train/memory_baselines_basic/run_ablation_prope_memory_two_chunk.sh.p0bak new file mode 100644 index 0000000000000000000000000000000000000000..84aae41e1c35e5f8235b5df62a8c59049412aa95 --- /dev/null +++ b/code/train/memory_baselines_basic/run_ablation_prope_memory_two_chunk.sh.p0bak @@ -0,0 +1,24 @@ +#!/bin/bash +# Ablation: PRoPE — CGLA (SSE-GLA) with camera-pose rotary position encoding +# (dfot `use_pose_rope` / PoseRoPE): q/k rotated by learned per-token camera-pose +# angles, so the GLA bilinear form depends on the RELATIVE camera pose. +# Same SSE-GLA backbone + contract + hyperparameters as run_ablation_cgla_memory_two_chunk.sh; +# only --cgla_mechanism prope differs (controlled ablation). +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +source "${SCRIPT_DIR}/common_sampling_two_chunk.sh" +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --context_source replay --prev_chunk_frames 81 --context_memory_frames 5 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_abl_prope_memory_two_chunk" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "abl_prope_memory_two_chunk" \ + --enable_context_memory --training_mode context --context_drop_prob 0.1 --cfg_target_only \ + --train_cam_pose \ + --use_moc --moc_temperature 1.0 \ + --use_cgla_memory --cgla_every_n_blocks 1 --cgla_aux_loss_weight 0.01 --cgla_mechanism prope \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" \ + "${SAMPLING_TWO_CHUNK_FLAGS[@]}" \ + 2>&1 | tee "${LOG_DIR}/abl_prope_memory_two_chunk_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/memory_baselines_basic/run_ablation_spatial_concat_text_two_chunk.sh b/code/train/memory_baselines_basic/run_ablation_spatial_concat_text_two_chunk.sh new file mode 100644 index 0000000000000000000000000000000000000000..c71ddcbe74e3315df5ca1d7b4ee05f658e57d85a --- /dev/null +++ b/code/train/memory_baselines_basic/run_ablation_spatial_concat_text_two_chunk.sh @@ -0,0 +1,21 @@ +#!/bin/bash +# Ablation: SpatialGridMemory + text KV concatenation read-out + 2-chunk monitor. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +source "${SCRIPT_DIR}/common_sampling_two_chunk.sh" +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --context_source prev_chunk_tail --context_memory_frames 1 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_abl_spatial_concat_text_two_chunk" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "abl_spatial_concat_text_two_chunk" \ + --enable_context_memory --training_mode context --condition_t2v_ratio 0.10 --condition_i2v_ratio 0.10 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --spike_threshold "${SPIKE_THRESHOLD:-15.0}" \ + --use_moc --moc_temperature 1.0 \ + --use_spatial_memory --spatial_memory_tokens 64 --spatial_memory_inject_mode concat_text \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" \ + "${SAMPLING_TWO_CHUNK_FLAGS[@]}" \ + 2>&1 | tee "${LOG_DIR}/abl_spatial_concat_text_two_chunk_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/memory_baselines_basic/run_ablation_spatial_cross_attn_readout_two_chunk.sh b/code/train/memory_baselines_basic/run_ablation_spatial_cross_attn_readout_two_chunk.sh new file mode 100644 index 0000000000000000000000000000000000000000..886e0921e4da8c4a221e548261750d252d5907fd --- /dev/null +++ b/code/train/memory_baselines_basic/run_ablation_spatial_cross_attn_readout_two_chunk.sh @@ -0,0 +1,21 @@ +#!/bin/bash +# Ablation: SpatialGridMemory + dedicated cross-attention read-out + 2-chunk monitor. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +source "${SCRIPT_DIR}/common_sampling_two_chunk.sh" +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --context_source prev_chunk_tail --context_memory_frames 1 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_abl_spatial_cross_attn_readout_two_chunk" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "abl_spatial_cross_attn_readout_two_chunk" \ + --enable_context_memory --training_mode context --condition_t2v_ratio 0.10 --condition_i2v_ratio 0.10 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --spike_threshold "${SPIKE_THRESHOLD:-15.0}" \ + --use_moc --moc_temperature 1.0 \ + --use_spatial_memory --spatial_memory_tokens 64 --spatial_memory_inject_mode cross_attn_readout \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" \ + "${SAMPLING_TWO_CHUNK_FLAGS[@]}" \ + 2>&1 | tee "${LOG_DIR}/abl_spatial_cross_attn_readout_two_chunk_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/memory_baselines_basic/run_ablation_spatial_inject_none_two_chunk.sh b/code/train/memory_baselines_basic/run_ablation_spatial_inject_none_two_chunk.sh new file mode 100644 index 0000000000000000000000000000000000000000..03002a76153ca90734b1535cbc63df418a81ff3e --- /dev/null +++ b/code/train/memory_baselines_basic/run_ablation_spatial_inject_none_two_chunk.sh @@ -0,0 +1,21 @@ +#!/bin/bash +# Ablation: SpatialGridMemory stored but not injected into the generator. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +source "${SCRIPT_DIR}/common_sampling_two_chunk.sh" +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --context_source prev_chunk_tail --context_memory_frames 1 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_abl_spatial_inject_none_two_chunk" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "abl_spatial_inject_none_two_chunk" \ + --enable_context_memory --training_mode context --condition_t2v_ratio 0.10 --condition_i2v_ratio 0.10 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --spike_threshold "${SPIKE_THRESHOLD:-15.0}" \ + --use_moc --moc_temperature 1.0 \ + --use_spatial_memory --spatial_memory_tokens 64 --spatial_memory_inject_mode none \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" \ + "${SAMPLING_TWO_CHUNK_FLAGS[@]}" \ + 2>&1 | tee "${LOG_DIR}/abl_spatial_inject_none_two_chunk_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/memory_baselines_basic/run_ablation_ucpe_memory_two_chunk.sh b/code/train/memory_baselines_basic/run_ablation_ucpe_memory_two_chunk.sh new file mode 100644 index 0000000000000000000000000000000000000000..8d2b1372e797eec1dd0114c229c0c61aca3f3fc1 --- /dev/null +++ b/code/train/memory_baselines_basic/run_ablation_ucpe_memory_two_chunk.sh @@ -0,0 +1,25 @@ +#!/bin/bash +# Ablation: UCPE — CGLA (SSE-GLA) with camera-pose rotary PE (use_pose_rope, +# the relative/PRoPE pillar) PLUS an absolute-orientation camera encoder +# (cam_encoder added to x, zero-init). Camera pose used as both relative and +# absolute positional encoding on the GLA memory pathway. +# Same SSE-GLA backbone + contract + hyperparameters as run_ablation_cgla_memory_two_chunk.sh; +# only --cgla_mechanism ucpe differs (controlled ablation). +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +source "${SCRIPT_DIR}/common_sampling_two_chunk.sh" +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --context_source replay --prev_chunk_frames 81 --context_memory_frames 5 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_abl_ucpe_memory_two_chunk" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "abl_ucpe_memory_two_chunk" \ + --enable_context_memory --training_mode context --context_drop_prob 0.1 --cfg_target_only \ + --train_cam_pose --add_action_attn \ + --use_moc --moc_temperature 1.0 \ + --use_cgla_memory --cgla_every_n_blocks 1 --cgla_aux_loss_weight 0.01 --cgla_mechanism ucpe \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" \ + "${SAMPLING_TWO_CHUNK_FLAGS[@]}" \ + 2>&1 | tee "${LOG_DIR}/abl_ucpe_memory_two_chunk_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/memory_baselines_basic/run_ablation_ucpe_memory_two_chunk.sh.p0bak b/code/train/memory_baselines_basic/run_ablation_ucpe_memory_two_chunk.sh.p0bak new file mode 100644 index 0000000000000000000000000000000000000000..67eacdc72ceeb3c51f64aef3d571da954f1ccada --- /dev/null +++ b/code/train/memory_baselines_basic/run_ablation_ucpe_memory_two_chunk.sh.p0bak @@ -0,0 +1,25 @@ +#!/bin/bash +# Ablation: UCPE — CGLA (SSE-GLA) with camera-pose rotary PE (use_pose_rope, +# the relative/PRoPE pillar) PLUS an absolute-orientation camera encoder +# (cam_encoder added to x, zero-init). Camera pose used as both relative and +# absolute positional encoding on the GLA memory pathway. +# Same SSE-GLA backbone + contract + hyperparameters as run_ablation_cgla_memory_two_chunk.sh; +# only --cgla_mechanism ucpe differs (controlled ablation). +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +source "${SCRIPT_DIR}/common_sampling_two_chunk.sh" +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --context_source replay --prev_chunk_frames 81 --context_memory_frames 5 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_abl_ucpe_memory_two_chunk" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "abl_ucpe_memory_two_chunk" \ + --enable_context_memory --training_mode context --context_drop_prob 0.1 --cfg_target_only \ + --train_cam_pose \ + --use_moc --moc_temperature 1.0 \ + --use_cgla_memory --cgla_every_n_blocks 1 --cgla_aux_loss_weight 0.01 --cgla_mechanism ucpe \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" \ + "${SAMPLING_TWO_CHUNK_FLAGS[@]}" \ + 2>&1 | tee "${LOG_DIR}/abl_ucpe_memory_two_chunk_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/memory_baselines_basic/run_ablation_videossm_hybrid_two_chunk.sh b/code/train/memory_baselines_basic/run_ablation_videossm_hybrid_two_chunk.sh new file mode 100644 index 0000000000000000000000000000000000000000..668ea2be7ce72340cd972d973eb8862cfd232276 --- /dev/null +++ b/code/train/memory_baselines_basic/run_ablation_videossm_hybrid_two_chunk.sh @@ -0,0 +1,21 @@ +#!/bin/bash +# Ablation: legacy VideoSSM hybrid + 2-chunk monitor. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +source "${SCRIPT_DIR}/common_sampling_two_chunk.sh" +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --enable_fov_retrieval --fov_top_k 4 --context_memory_frames 5 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_abl_videossm_hybrid_two_chunk" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "abl_videossm_hybrid_two_chunk" \ + --enable_context_memory --training_mode context --condition_t2v_ratio 0.10 --condition_i2v_ratio 0.10 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --spike_threshold "${SPIKE_THRESHOLD:-15.0}" \ + --use_moc --moc_temperature 1.0 \ + --use_videossm_hybrid --videossm_every_n_blocks 4 --videossm_kernel_size 3 --videossm_expand 2 \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" \ + "${SAMPLING_TWO_CHUNK_FLAGS[@]}" \ + 2>&1 | tee "${LOG_DIR}/abl_videossm_hybrid_two_chunk_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/memory_baselines_basic/run_all_ablations_two_chunk.sh b/code/train/memory_baselines_basic/run_all_ablations_two_chunk.sh new file mode 100644 index 0000000000000000000000000000000000000000..4b066fbe18cd6a7ab5a6e64acf3097586d4d0703 --- /dev/null +++ b/code/train/memory_baselines_basic/run_all_ablations_two_chunk.sh @@ -0,0 +1,16 @@ +#!/bin/bash +# Convenience launcher for the paper ablation matrix. This is long-running. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +bash "${SCRIPT_DIR}/run_ablation_no_memory_baseline_two_chunk.sh" +bash "${SCRIPT_DIR}/run_ablation_framepack_weight_two_chunk.sh" +bash "${SCRIPT_DIR}/run_ablation_framepack_len_r2_two_chunk.sh" +bash "${SCRIPT_DIR}/run_ablation_framepack_len_r4_two_chunk.sh" +bash "${SCRIPT_DIR}/run_ablation_framepack_hybrid_r2_weight_two_chunk.sh" +bash "${SCRIPT_DIR}/run_ablation_framepack_hybrid_r4_weight_two_chunk.sh" +bash "${SCRIPT_DIR}/run_ablation_spatial_inject_none_two_chunk.sh" +bash "${SCRIPT_DIR}/run_ablation_spatial_concat_text_two_chunk.sh" +bash "${SCRIPT_DIR}/run_ablation_spatial_cross_attn_readout_two_chunk.sh" +bash "${SCRIPT_DIR}/run_ablation_videossm_hybrid_two_chunk.sh" +bash "${SCRIPT_DIR}/run_ablation_block_wise_ssm_two_chunk.sh" diff --git a/code/train/memory_baselines_basic/run_framepack_baseline.sh b/code/train/memory_baselines_basic/run_framepack_baseline.sh new file mode 100644 index 0000000000000000000000000000000000000000..bcf41fd2bfc05fc2e9082c4d2e8c47e946a864c6 --- /dev/null +++ b/code/train/memory_baselines_basic/run_framepack_baseline.sh @@ -0,0 +1,22 @@ +#!/bin/bash +# FramePack/FAR baseline: weight-only (no length compression) +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +export CONTEXT_POSITION=prefix +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --context_source prev_chunk_tail --context_memory_frames 81 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_framepack_weight_only" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "memory_baseline_framepack_weight_only" \ + --enable_context_memory --training_mode context --context_drop_prob 0.1 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --use_moc --moc_temperature 1.0 \ + --use_framepack_memory --context_temporal_decay 0.9 --context_attention_weight 1.0 \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" --enable_video_sampling --sampling_atomic_left_right --sampling_interval_steps "${SAMPLING_INTERVAL_STEPS:-1000}" \ + --sampling_num_inference_steps "${SAMPLING_NUM_INFERENCE_STEPS:-50}" --sampling_negative_prompt "oversaturated colors, overexposed, static, blurry details" \ + --sampling_height "${SAMPLING_HEIGHT:-352}" --sampling_width "${SAMPLING_WIDTH:-640}" --sampling_num_frames "${SAMPLING_NUM_FRAMES:-81}" --samples_per_epoch 0 --sampling_action_path "${sampling_action_path}" \ + 2>&1 | tee "${LOG_DIR}/memory_baseline_framepack_weight_only_$(date +%Y%m%d_%H%M%S).log" + diff --git a/code/train/memory_baselines_basic/run_framepack_lencompress_r2.sh b/code/train/memory_baselines_basic/run_framepack_lencompress_r2.sh new file mode 100644 index 0000000000000000000000000000000000000000..490052ebc3404a9cba6202696cf5629219eb844d --- /dev/null +++ b/code/train/memory_baselines_basic/run_framepack_lencompress_r2.sh @@ -0,0 +1,25 @@ +#!/bin/bash +# FramePack/FAR baseline: frame-level length compression K->K' with ratio=2 +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +export CONTEXT_POSITION=suffix + +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --context_source prev_chunk_tail --context_memory_frames 81 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_framepack_lencompress_r2" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "memory_baseline_framepack_lencompress_r2" \ + --enable_context_memory --training_mode context --condition_t2v_ratio 0.10 --condition_i2v_ratio 0.10 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --spike_threshold "${SPIKE_THRESHOLD:-15.0}" \ + --use_moc --moc_temperature 1.0 \ + --use_framepack_length_compress --framepack_ratio 2 \ + --framepack_length_strategy packed_multiscale --framepack_multiscale_w2 0.25 --framepack_multiscale_w4 0.15 \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" --enable_video_sampling --sampling_atomic_left_right --sampling_interval_steps "${SAMPLING_INTERVAL_STEPS:-1000}" \ + --sampling_num_inference_steps "${SAMPLING_NUM_INFERENCE_STEPS:-50}" --sampling_negative_prompt "oversaturated colors, overexposed, static, blurry details" \ + --sampling_height "${SAMPLING_HEIGHT:-352}" --sampling_width "${SAMPLING_WIDTH:-640}" --sampling_num_frames "${SAMPLING_NUM_FRAMES:-81}" --samples_per_epoch 0 --sampling_action_path "${sampling_action_path}" \ + 2>&1 | tee "${LOG_DIR}/memory_baseline_framepack_lencompress_r2_$(date +%Y%m%d_%H%M%S).log" + diff --git a/code/train/memory_baselines_basic/run_framepack_lencompress_r4.sh b/code/train/memory_baselines_basic/run_framepack_lencompress_r4.sh new file mode 100644 index 0000000000000000000000000000000000000000..0d3a929561ec7d508c425d42a22da5feb3aec5f4 --- /dev/null +++ b/code/train/memory_baselines_basic/run_framepack_lencompress_r4.sh @@ -0,0 +1,24 @@ +#!/bin/bash +# FramePack/FAR baseline: frame-level length compression K->K' with ratio=4. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +export CONTEXT_POSITION=suffix + +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --context_source prev_chunk_tail --context_memory_frames 81 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_framepack_lencompress_r4" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "memory_baseline_framepack_lencompress_r4" \ + --enable_context_memory --training_mode context --condition_t2v_ratio 0.10 --condition_i2v_ratio 0.10 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --spike_threshold "${SPIKE_THRESHOLD:-15.0}" \ + --use_moc --moc_temperature 1.0 \ + --use_framepack_length_compress --framepack_ratio 4 \ + --framepack_length_strategy packed_multiscale --framepack_multiscale_w2 0.25 --framepack_multiscale_w4 0.15 \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" --enable_video_sampling --sampling_atomic_left_right --sampling_interval_steps "${SAMPLING_INTERVAL_STEPS:-1000}" \ + --sampling_num_inference_steps "${SAMPLING_NUM_INFERENCE_STEPS:-50}" --sampling_negative_prompt "oversaturated colors, overexposed, static, blurry details" \ + --sampling_height "${SAMPLING_HEIGHT:-352}" --sampling_width "${SAMPLING_WIDTH:-640}" --sampling_num_frames "${SAMPLING_NUM_FRAMES:-81}" --samples_per_epoch 0 --sampling_action_path "${sampling_action_path}" \ + 2>&1 | tee "${LOG_DIR}/memory_baseline_framepack_lencompress_r4_$(date +%Y%m%d_%H%M%S).log" diff --git a/code/train/memory_baselines_basic/run_geometry_spatial_memory_baseline.sh b/code/train/memory_baselines_basic/run_geometry_spatial_memory_baseline.sh new file mode 100644 index 0000000000000000000000000000000000000000..c4f4fa491eea4ed9dd51fbc8c2df95188db627e2 --- /dev/null +++ b/code/train/memory_baselines_basic/run_geometry_spatial_memory_baseline.sh @@ -0,0 +1,33 @@ +#!/bin/bash +# Geometry-grounded Spatial Memory. +# Requires metadata column GEOMETRY_MEMORY_COLUMN whose values point to +# TSDF/point-cloud-rendered static videos (for example Vid_masktarget.mp4). +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" + +GEOMETRY_MEMORY_COLUMN="${GEOMETRY_MEMORY_COLUMN:-geometry_memory}" +GEOMETRY_MEMORY_ROOT="${GEOMETRY_MEMORY_ROOT:-${dataset_base_path}}" +EXTRA_ARGS=() +[ -n "${MAX_TRAIN_STEPS:-}" ] && EXTRA_ARGS+=(--max_train_steps "${MAX_TRAIN_STEPS}") +[ -n "${PROGRESS_TOTAL_STEPS:-}" ] && EXTRA_ARGS+=(--progress_total_steps "${PROGRESS_TOTAL_STEPS}") + +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --geometry_memory_column "${GEOMETRY_MEMORY_COLUMN}" --geometry_memory_root "${GEOMETRY_MEMORY_ROOT}" \ + --context_source prev_chunk_tail --context_memory_frames 1 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers "${NUM_WORKERS:-16}" \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_geometry_spatial_mem" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "memory_baseline_geometry_spatial_mem" \ + --enable_context_memory --training_mode context --condition_t2v_ratio 0.10 --condition_i2v_ratio 0.10 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --spike_threshold "${SPIKE_THRESHOLD:-15.0}" \ + --use_moc --moc_temperature 1.0 \ + --use_geometry_spatial_memory --geometry_spatial_memory_tokens 64 \ + --geometry_spatial_memory_grid 8 --geometry_spatial_memory_temporal_bins 4 \ + --geometry_spatial_memory_inject_mode "${GEOMETRY_MEMORY_INJECT_MODE:-concat_text}" \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" \ + "${EXTRA_ARGS[@]}" \ + 2>&1 | tee "${LOG_DIR}/memory_baseline_geometry_spatial_mem_$(date +%Y%m%d_%H%M%S).log" + diff --git a/code/train/memory_baselines_basic/run_spatial_memory_baseline.sh b/code/train/memory_baselines_basic/run_spatial_memory_baseline.sh new file mode 100644 index 0000000000000000000000000000000000000000..26d698d4123d330edf6f5886de097174c86437fa --- /dev/null +++ b/code/train/memory_baselines_basic/run_spatial_memory_baseline.sh @@ -0,0 +1,24 @@ +#!/bin/bash +# Token-grid baseline (historically named spatial_mem): append pooled context +# tokens to cross-attn. This is not the depth/TSDF geometry-grounded method. +# Optional: --spatial_memory_inject_mode {concat_text,none,cross_attn_readout}; optional train monitor --sampling_two_chunk_memory (see README). +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --context_source prev_chunk_tail --context_memory_frames 1 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_spatial_mem" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "memory_baseline_spatial_mem" \ + --enable_context_memory --training_mode context --condition_t2v_ratio 0.10 --condition_i2v_ratio 0.10 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --spike_threshold "${SPIKE_THRESHOLD:-15.0}" \ + --use_moc --moc_temperature 1.0 \ + --use_spatial_memory --spatial_memory_tokens 64 \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" --enable_video_sampling --sampling_atomic_left_right --sampling_interval_steps "${SAMPLING_INTERVAL_STEPS:-1000}" \ + --sampling_num_inference_steps "${SAMPLING_NUM_INFERENCE_STEPS:-50}" --sampling_negative_prompt "oversaturated colors, overexposed, static, blurry details" \ + --sampling_height "${SAMPLING_HEIGHT:-352}" --sampling_width "${SAMPLING_WIDTH:-640}" --sampling_num_frames "${SAMPLING_NUM_FRAMES:-81}" --samples_per_epoch 0 --sampling_action_path "${sampling_action_path}" \ + 2>&1 | tee "${LOG_DIR}/memory_baseline_spatial_mem_$(date +%Y%m%d_%H%M%S).log" + diff --git a/code/train/memory_baselines_basic/run_videossm_hybrid_baseline.sh b/code/train/memory_baselines_basic/run_videossm_hybrid_baseline.sh new file mode 100644 index 0000000000000000000000000000000000000000..fa9a307ba5354b65fa1a60d55e297fc04e158f3e --- /dev/null +++ b/code/train/memory_baselines_basic/run_videossm_hybrid_baseline.sh @@ -0,0 +1,23 @@ +#!/bin/bash +# Legacy VideoSSM hybrid (depthwise conv). For paper-aligned SSM narrative prefer --use_block_wise_ssm in a separate script. +# VideoSSM baseline: hybrid state-space memory in DiT blocks +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "${SCRIPT_DIR}/common_env.sh" +accelerate launch src/model_training/train.py \ + --dataset_base_path "${dataset_base_path}" --dataset_metadata_path "${dataset_base_path}/${METADATA_NAME}" \ + --enable_fov_retrieval --fov_top_k 4 --context_memory_frames 5 --use_rt_relative --height 352 --width 640 \ + --dataset_repeat 1 --per_device_train_batch_size 1 --gradient_accumulation_steps 1 --num_workers 16 \ + --model_paths "${model_paths}" --learning_rate 5e-5 --num_epochs 1 --remove_prefix_in_ckpt "${remove_prefix_in_ckpt}" \ + --output_path "${output_base}_videossm_hybrid" --trainable_models dit --ckpt_interval "${CKPT_INTERVAL:-1000}" --save_full_model \ + --wandb_run_name "memory_baseline_videossm_hybrid" \ + --enable_context_memory --training_mode context --condition_t2v_ratio 0.10 --condition_i2v_ratio 0.10 --cfg_target_only \ + --train_cam_pose --add_action_attn --action_use_temporal_attention \ + --spike_threshold "${SPIKE_THRESHOLD:-15.0}" \ + --use_moc --moc_temperature 1.0 \ + --use_videossm_hybrid --videossm_every_n_blocks 4 --videossm_kernel_size 3 --videossm_expand 2 \ + --timestep_shift "${TIMESTEP_SHIFT:-15}" --enable_video_sampling --sampling_atomic_left_right --sampling_interval_steps "${SAMPLING_INTERVAL_STEPS:-1000}" \ + --sampling_num_inference_steps "${SAMPLING_NUM_INFERENCE_STEPS:-50}" --sampling_negative_prompt "oversaturated colors, overexposed, static, blurry details" \ + --sampling_height "${SAMPLING_HEIGHT:-352}" --sampling_width "${SAMPLING_WIDTH:-640}" --sampling_num_frames "${SAMPLING_NUM_FRAMES:-81}" --samples_per_epoch 0 --sampling_action_path "${sampling_action_path}" \ + 2>&1 | tee "${LOG_DIR}/memory_baseline_videossm_hybrid_$(date +%Y%m%d_%H%M%S).log" +