[ci]: collect the tracker unit tests in the unit lane

fastvideo/tests/training/test_trackers.py (CPU-only, fake-swanlab
fixtures) was not collected by any lane - the training lanes pick
specific subdirs and run_unit_test skipped tests/training entirely.

docker/Dockerfile

@@ -68,7 +68,7 @@ ARG FLASH_ATTN_WHEEL_RELEASE_ARM64=https://github.com/mjun0812/flash-attention-p
 # cutlass-4.4 `cute.core.ThrMma` API, which crashes on the cutlass-dsl 4.5 that
 # flashinfer/quack pull in. After the wheel install we overlay this cutlass-4.5-safe
 # upstream cute (flash-attn-4) so the image runs FA4 instead of the FA2 fallback.
-ARG FA4_CUTE_REF=940cd9680f3315f2f06b43ab5bea2c2cf2d96806
+ARG FA4_CUTE_REF=82d6441eec5d4dfec120153db2c0145ae855a083
 
 # Provided automatically by BuildKit/buildx (e.g. "amd64" / "arm64") and used to
 # select the prebuilt flash-attn wheel. Empty under a plain `docker build` without
@@ -161,7 +161,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
         uv pip install flash-attn==${FLASH_ATTN_VERSION} --no-build-isolation; \
     fi
 
-# Overlay the cutlass-4.5-safe upstream FA4 cute (FA4_CUTE_REF) over the
+# Overlay the CuTe-DSL-4.6-compatible upstream FA4 cute (FA4_CUTE_REF) over the
 # wheel/source one so the image runs FA4, not the FA2 fallback. This pulls the FA4
 # runtime stack (cutlass-dsl, quack-kernels, apache-tvm-ffi, torch-c-dlpack-ext) --
 # the same deps the [dreamverse] extra already installs in CI; the installed torch

docs/training/trackers.md

@@ -0,0 +1,91 @@
+# Training Trackers
+
+FastVideo can send training metrics and validation media to Weights & Biases
+or SwanLab. Tracking runs only on global rank 0, and local tracker files are
+stored under `<output_dir>/tracker`.
+
+## Supported Trackers
+
+| Value | Backend | Installation |
+|-------|---------|--------------|
+| `wandb` | Weights & Biases | Included with FastVideo |
+| `swanlab` | SwanLab | Install the optional `swanlab` dependency |
+| `none` | Disable external tracking | No additional package |
+
+You can enable more than one backend, for example `trackers: [wandb, swanlab]`.
+Metrics and validation media are converted to the artifact type required by
+each backend.
+
+## Install SwanLab
+
+For a published FastVideo installation, install the SwanLab extra:
+
+```bash
+uv pip install "fastvideo[swanlab]"
+```
+
+For an editable source checkout, include the same extra during installation:
+
+```bash
+uv pip install -e ".[swanlab]"
+```
+
+If FastVideo is already installed, you can install the compatible SDK directly:
+
+```bash
+uv pip install "swanlab>=0.6.7"
+```
+
+Authenticate once before starting a training run:
+
+```bash
+swanlab login
+```
+
+See the [SwanLab login documentation](https://docs.swanlab.cn/en/api/cli-swanlab-login.html)
+for non-interactive and self-hosted setups.
+
+## Configure Tracking
+
+Select SwanLab in the YAML config used by the modular training framework:
+
+```yaml
+training:
+  checkpoint:
+    output_dir: outputs/my_run
+  tracker:
+    trackers: [swanlab]
+    project_name: my_project
+    run_name: my_run
+```
+
+To log to both supported services:
+
+```yaml
+training:
+  tracker:
+    trackers: [wandb, swanlab]
+    project_name: my_project
+    run_name: my_run
+```
+
+An empty or omitted `trackers` list selects W&B when `project_name` is set.
+Use an explicit `none` entry to disable external tracking:
+
+```yaml
+training:
+  tracker:
+    trackers: [none]
+```
+
+## Validation Videos
+
+SwanLab currently accepts GIF video artifacts. FastVideo converts validation
+MP4 files and in-memory video arrays to GIF automatically before logging them.
+For video files, FastVideo uses the sampling frame rate supplied by the caller,
+or the source file's frame rate when no value is supplied. In-memory arrays use
+the frame rate supplied by the caller. Both forms fall back to 16 FPS when no
+frame rate is available.
+
+For details about configuring validation callbacks, see
+[Training Infrastructure](train_infra.md#callbacks-pluggable-hooks).

docs/training/train_infra.md

@@ -161,6 +161,9 @@ training:
     decay_interval_steps: 0
 ```
 
+See [Training Trackers](trackers.md) to configure Weights & Biases or SwanLab,
+including SwanLab installation and authentication.
+
 ### `callbacks` — Pluggable hooks
 
 Callbacks run at specific points in the training loop (before/after optimizer

fastvideo/distributed/parallel_state.py

@@ -23,6 +23,7 @@ If you only need to use the distributed environment without model parallelism,
  you can skip the model parallel initialization and destruction steps.
 """
 import contextlib
+import gc
 import os
 import pickle
 import weakref
@@ -1006,6 +1007,13 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
     if shutdown_ray:
         import ray  # Lazy import Ray
         ray.shutdown()
+    # Actually free GPU memory, as the name promises. FSDP-wrapped modules sit
+    # in reference cycles (module <-> FSDP state <-> hooks), so dropping the
+    # last user reference does not free their parameters until a gc pass runs.
+    # Without this, back-to-back model-load tests in one pytest session OOM.
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
 
 
 def get_same_node_ranks(pg: ProcessGroup | StatelessProcessGroup, source_rank: int = 0) -> list[int]:

fastvideo/tests/api/test_compat_translation.py

@@ -17,46 +17,6 @@ class TestLegacyTorchCompileKwargsTranslation:
     first-class :class:`CompileConfig` fields and anything unknown falls
     into ``extras``."""
 
-    def test_all_typed_keys_promoted(self) -> None:
-        config = legacy_from_pretrained_to_config(
-            "/models/ltx2",
-            {
-                "enable_torch_compile": True,
-                "torch_compile_kwargs": {
-                    "backend": "inductor",
-                    "fullgraph": True,
-                    "mode": "max-autotune-no-cudagraphs",
-                    "dynamic": False,
-                },
-            },
-        )
-        compile_config = config.engine.compile
-        assert compile_config.enabled is True
-        assert compile_config.backend == "inductor"
-        assert compile_config.fullgraph is True
-        assert compile_config.mode == "max-autotune-no-cudagraphs"
-        assert compile_config.dynamic is False
-        assert compile_config.extras == {}
-
-    def test_unknown_keys_land_in_extras(self) -> None:
-        config = legacy_from_pretrained_to_config(
-            "/models/ltx2",
-            {
-                "enable_torch_compile": True,
-                "torch_compile_kwargs": {
-                    "backend": "inductor",
-                    "options": {"triton.cudagraphs": False},
-                    "disable": False,
-                },
-            },
-        )
-        compile_config = config.engine.compile
-        assert compile_config.backend == "inductor"
-        assert compile_config.extras == {
-            "options": {"triton.cudagraphs": False},
-            "disable": False,
-        }
-
     def test_empty_kwargs_produces_empty_extras(self) -> None:
         config = legacy_from_pretrained_to_config(
             "/models/ltx2",
@@ -117,13 +77,6 @@ class TestLegacyLtx2VaeTilingTranslation:
     ``generator.pipeline.vae_tiling``; reverse direction emits the
     legacy name back to FastVideoArgs."""
 
-    def test_forward_routes_to_pipeline_vae_tiling(self) -> None:
-        config = legacy_from_pretrained_to_config(
-            "/models/ltx2",
-            {"ltx2_vae_tiling": False},
-        )
-        assert config.pipeline.vae_tiling is False
-
     def test_true_round_trips(self) -> None:
         config = legacy_from_pretrained_to_config(
             "/models/ltx2",

fastvideo/tests/contract/test_ci_test_collection.py

@@ -35,7 +35,6 @@ ALLOWLIST = {
     "hooks": "no lane yet — run manually",
     "layers": "no lane yet — torchrun FSDP dispatch tests, run manually",
     "nightly": "by design: nightly cadence, not per-PR",
-    "ops": "no lane yet — GPU op tests, run manually",
     "modal": "CI infrastructure itself, not a test suite",
 }
 

fastvideo/tests/modal/pr_test.py

@@ -187,6 +187,8 @@ def run_transformer_tests():
 
 
 @app.function(gpu="L40S:4",
+              cpu=8.0,
+              memory=32768,
               image=image,
               timeout=900,
               secrets=[
@@ -201,6 +203,8 @@ def run_training_tests():
 
 
 @app.function(gpu="L40S:2",
+              cpu=8.0,
+              memory=32768,
               image=image,
               timeout=900,
               secrets=[
@@ -276,7 +280,7 @@ def run_self_forcing_tests():
 @app.function(gpu="L40S:1", image=image, timeout=900)
 def run_unit_test():
     run_test(
-        "pytest ./fastvideo/tests/api/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ ./fastvideo/tests/stages/ --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py --ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs"
+        "pytest ./fastvideo/tests/api/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ ./fastvideo/tests/stages/ ./fastvideo/tests/ops/ ./fastvideo/tests/training/test_trackers.py --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py --ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs"
     )
 
 
@@ -318,6 +322,8 @@ def run_dreamverse_app_tests():
 
 
 @app.function(gpu="L40S:1",
+              cpu=8.0,
+              memory=32768,
               image=image,
               timeout=1800,
               secrets=[
@@ -400,6 +406,8 @@ def run_lora_extraction_tests():
 
 
 @app.function(gpu="L40S:2",
+              cpu=8.0,
+              memory=32768,
               image=image,
               timeout=1800,
               secrets=[

fastvideo/tests/ops/quantization/test_absmax_fp8.py

@@ -34,6 +34,9 @@ class TestAbsMaxFP8LinearMethod(unittest.TestCase):
             method._convert_scale(scale)
 
     def test_create_weights_rejects_invalid_dtype(self):
+        # float64 is outside the supported set (bf16/fp16/fp32). The test
+        # originally passed float32, which create_weights has accepted since
+        # PR #981 — it never ran in CI, so the mismatch went unnoticed.
         method = AbsMaxFP8LinearMethod()
         layer = nn.Module()
         with self.assertRaisesRegex(AssertionError, "only supports"):
@@ -43,7 +46,7 @@ class TestAbsMaxFP8LinearMethod(unittest.TestCase):
                 output_partition_sizes=[3],
                 input_size=2,
                 output_size=3,
-                params_dtype=torch.float32,
+                params_dtype=torch.float64,
             )
 
     def test_absmax_fp8_parameter_weight_loader(self):

fastvideo/tests/training/test_trackers.py

@@ -0,0 +1,225 @@
+from __future__ import annotations
+
+import builtins
+from io import BytesIO
+from pathlib import Path
+import sys
+from types import ModuleType
+from typing import Any
+
+import numpy as np
+from PIL import Image
+import pytest
+
+from fastvideo.training import trackers as trackers_module
+from fastvideo.training.trackers import (
+    _coerce_video_fps,
+    _prepare_video_array,
+    BaseTracker,
+    SequentialTracker,
+    SwanlabTracker,
+    WandbTracker,
+)
+
+
+class _FakeSwanlab(ModuleType):
+    def __init__(self) -> None:
+        super().__init__("swanlab")
+        self.init_kwargs: dict[str, Any] | None = None
+        self.logged: list[tuple[dict[str, Any], int]] = []
+        self.video_paths: list[str] = []
+        self.finished = False
+
+    def init(self, **kwargs: Any) -> object:
+        self.init_kwargs = kwargs
+        return object()
+
+    def log(self, metrics: dict[str, Any], step: int) -> None:
+        self.logged.append((metrics, step))
+
+    def finish(self) -> None:
+        self.finished = True
+
+    def Video(self, file_path: str, caption: str | None = None) -> dict[str, Any]:  # noqa: N802
+        self.video_paths.append(file_path)
+        return {
+            "content": Path(file_path).read_bytes(),
+            "caption": caption,
+        }
+
+
+def _install_fake_swanlab(monkeypatch: pytest.MonkeyPatch) -> _FakeSwanlab:
+    swanlab = _FakeSwanlab()
+    monkeypatch.setitem(sys.modules, "swanlab", swanlab)
+    return swanlab
+
+
+def test_prepare_video_array_converts_tchw_to_thwc() -> None:
+    video = np.arange(2 * 3 * 2 * 4, dtype=np.uint8).reshape(2, 3, 2, 4)
+
+    prepared = _prepare_video_array(video)
+
+    np.testing.assert_array_equal(prepared, video.transpose(0, 2, 3, 1))
+
+
+def test_default_video_fps_is_16() -> None:
+    assert _coerce_video_fps(None) == 16
+
+
+def test_wandb_video_uses_shared_default_fps(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+) -> None:
+    wandb = ModuleType("wandb")
+    video_calls: list[dict[str, Any]] = []
+    wandb.init = lambda **_: object()
+
+    def video(data: Any, **kwargs: Any) -> dict[str, Any]:
+        video_calls.append(kwargs)
+        return {"data": data, **kwargs}
+
+    wandb.Video = video
+    monkeypatch.setitem(sys.modules, "wandb", wandb)
+    tracker = WandbTracker("project", str(tmp_path))
+
+    tracker.video(np.zeros((2, 3, 2, 2), dtype=np.uint8))
+
+    assert video_calls == [{"fps": 16, "format": "mp4"}]
+
+
+def test_prepare_video_array_tiles_and_pads_batches() -> None:
+    video = np.zeros((3, 1, 3, 1, 1), dtype=np.uint8)
+    video[0] = 1
+    video[1] = 2
+    video[2] = 3
+
+    prepared = _prepare_video_array(video)
+
+    assert prepared.shape == (1, 1, 4, 3)
+    np.testing.assert_array_equal(
+        prepared[0, 0],
+        np.array(
+            [
+                [1, 1, 1],
+                [2, 2, 2],
+                [3, 3, 3],
+                [0, 0, 0],
+            ],
+            dtype=np.uint8,
+        ),
+    )
+
+
+def test_swanlab_video_converts_array_to_gif_and_logs_artifact(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+) -> None:
+    swanlab = _install_fake_swanlab(monkeypatch)
+    tracker = SwanlabTracker("project", str(tmp_path), run_name="run")
+    video = np.zeros((2, 3, 2, 2), dtype=np.uint8)
+    video[0, 0] = 255
+    video[1, 2] = 255
+
+    artifact = tracker.video(video, caption="validation sample", fps=5, format="mp4")
+    tracker.log_artifacts({"validation/videos": [artifact]}, step=7)
+
+    assert artifact["content"].startswith((b"GIF87a", b"GIF89a"))
+    assert artifact["caption"] == "validation sample"
+    assert not Path(swanlab.video_paths[0]).exists()
+    with Image.open(BytesIO(artifact["content"])) as image:
+        assert image.n_frames == 2
+        assert image.info["duration"] == 200
+    assert swanlab.logged == [({"validation/videos": [artifact]}, 7)]
+
+
+@pytest.mark.parametrize(("requested_fps", "expected_fps"), [(None, 12.0), (2, 2.0)])
+def test_swanlab_video_uses_requested_or_source_fps_during_gif_conversion(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+    requested_fps: int | None,
+    expected_fps: float,
+) -> None:
+    _install_fake_swanlab(monkeypatch)
+    tracker = SwanlabTracker("project", str(tmp_path))
+    frames = [np.zeros((2, 2, 3), dtype=np.uint8), np.ones((2, 2, 3), dtype=np.uint8)]
+    written_fps: list[float] = []
+    write_gif = trackers_module._write_gif
+
+    monkeypatch.setattr(trackers_module, "_read_video_file", lambda _: (frames, 12.0))
+
+    def capture_fps(file_path: str, video_frames: Any, fps: float) -> None:
+        written_fps.append(fps)
+        write_gif(file_path, video_frames, fps)
+
+    monkeypatch.setattr(trackers_module, "_write_gif", capture_fps)
+
+    artifact = tracker.video(tmp_path / "validation.mp4", fps=requested_fps)
+
+    assert artifact["content"].startswith((b"GIF87a", b"GIF89a"))
+    assert written_fps == [expected_fps]
+
+
+def test_swanlab_video_passes_existing_gif_through(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+) -> None:
+    swanlab = _install_fake_swanlab(monkeypatch)
+    tracker = SwanlabTracker("project", str(tmp_path))
+    gif_path = tmp_path / "validation.gif"
+    Image.new("RGB", (2, 2)).save(gif_path, format="GIF")
+
+    artifact = tracker.video(gif_path, caption="existing")
+
+    assert swanlab.video_paths == [str(gif_path)]
+    assert gif_path.exists()
+    assert artifact["caption"] == "existing"
+
+
+def test_swanlab_missing_dependency_error_is_actionable(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+) -> None:
+    monkeypatch.delitem(sys.modules, "swanlab", raising=False)
+    real_import = builtins.__import__
+
+    def missing_swanlab(name: str, *args: Any, **kwargs: Any) -> Any:
+        if name == "swanlab":
+            raise ModuleNotFoundError("No module named 'swanlab'", name="swanlab")
+        return real_import(name, *args, **kwargs)
+
+    monkeypatch.setattr(builtins, "__import__", missing_swanlab)
+
+    with pytest.raises(ModuleNotFoundError, match=r"fastvideo\[swanlab\]"):
+        SwanlabTracker("project", str(tmp_path))
+
+
+class _RecordingTracker(BaseTracker):
+    def __init__(self, name: str) -> None:
+        super().__init__()
+        self.name = name
+        self.artifact_logs: list[dict[str, Any]] = []
+
+    def video(
+        self,
+        data: Any,
+        *,
+        caption: str | None = None,
+        fps: int | None = None,
+        format: str | None = None,
+    ) -> str:
+        return f"{self.name}:{data}"
+
+    def log_artifacts(self, artifacts: dict[str, Any], step: int) -> None:
+        self.artifact_logs.append(artifacts)
+
+
+def test_sequential_tracker_logs_backend_specific_video_artifacts() -> None:
+    first = _RecordingTracker("first")
+    second = _RecordingTracker("second")
+    tracker = SequentialTracker([first, second])
+
+    artifact = tracker.video("sample.mp4")
+    tracker.log_artifacts({"validation/videos": [artifact]}, step=3)
+
+    assert first.artifact_logs == [{"validation/videos": ["first:sample.mp4"]}]
+    assert second.artifact_logs == [{"validation/videos": ["second:sample.mp4"]}]

fastvideo/tests/transformers/test_hunyuanvideo.py

@@ -1,104 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-import glob
-import json
-import os
-import pytest
-import torch
-
-from fastvideo.configs.pipelines.base import PipelineConfig
-from fastvideo.distributed.parallel_state import get_sp_parallel_rank, get_sp_world_size
-from fastvideo.logger import init_logger
-from fastvideo.models.loader.component_loader import TransformerLoader
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.utils import maybe_download_model
-from fastvideo.forward_context import set_forward_context
-from fastvideo.configs.models.dits import HunyuanVideoConfig
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-
-logger = init_logger(__name__)
-
-os.environ["MASTER_ADDR"] = "localhost"
-os.environ["MASTER_PORT"] = "29503"
-
-BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH))
-TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
-CONFIG_PATH = os.path.join(TRANSFORMER_PATH, "config.json")
-
-LOCAL_RANK = 0
-RANK = 0
-WORLD_SIZE = 1
-
-# Latent generated on commit c021e8a27cf437ac22827f2bc58b7f006561317f with 1 x L40S
-REFERENCE_LATENT = 93.23263549804688
-
-
-@pytest.mark.usefixtures("distributed_setup")
-@pytest.mark.skip(reason="Skipping test_hunyuanvideo_distributed")
-def test_hunyuanvideo_distributed():
-    logger.info(f"Initializing process: rank={RANK}, local_rank={LOCAL_RANK}, world_size={WORLD_SIZE}")
-
-    torch.cuda.set_device(f"cuda:{LOCAL_RANK}")
-
-    # Get tensor parallel info
-    sp_rank = get_sp_parallel_rank()
-    sp_world_size = get_sp_world_size()
-
-    logger.info(f"Process rank {RANK} initialized with SP rank {sp_rank} in SP world size {sp_world_size}")
-
-    config = json.load(open(CONFIG_PATH))
-    # remove   "_class_name": "HunyuanVideoTransformer3DModel",   "_diffusers_version": "0.32.0.dev0",
-    # TODO: write normalize config function
-    config.pop("_class_name")
-    config.pop("_diffusers_version")
-
-    precision_str = "bf16"
-    args = FastVideoArgs(
-        model_path=TRANSFORMER_PATH,
-        dit_cpu_offload=True,
-        pipeline_config=PipelineConfig(dit_config=HunyuanVideoConfig(), dit_precision=precision_str),
-    )
-    args.device = torch.device(f"cuda:{LOCAL_RANK}")
-
-    loader = TransformerLoader()
-    model = loader.load(TRANSFORMER_PATH, args)
-
-    model.eval()
-
-    # Create random inputs for testing
-    batch_size = 1
-    seq_len = 3
-    device = torch.device(f"cuda:{LOCAL_RANK}")
-
-    # Video latents [B, C, T, H, W]
-    hidden_states = torch.randn(batch_size, 16, 8, 16, 16, device=device, dtype=torch.bfloat16)
-    chunk_per_rank = hidden_states.shape[2] // sp_world_size
-    hidden_states = hidden_states[:, :, sp_rank * chunk_per_rank : (sp_rank + 1) * chunk_per_rank]
-
-    # Text embeddings [B, L, D] (including global token)
-    encoder_hidden_states = torch.randn(batch_size, seq_len + 1, 4096, device=device, dtype=torch.bfloat16)
-
-    # Timestep
-    timestep = torch.tensor([500], device=device, dtype=torch.bfloat16)
-    forward_batch = ForwardBatch(
-        data_type="dummy",
-    )
-
-    # Disable gradients for inference
-    with torch.no_grad():
-        # Run inference on model
-        with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
-            with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=forward_batch):
-                output = model(
-                    hidden_states=hidden_states,
-                    encoder_hidden_states=encoder_hidden_states,
-                    timestep=timestep,
-                )
-
-    latent = output.double().sum().item()
-
-    # Check if latents are similar
-    diff_output_latents = abs(REFERENCE_LATENT - latent)
-    logger.info(f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}")
-    assert diff_output_latents < 1e-4, f"Output latents differ significantly: max diff = {diff_output_latents}"

fastvideo/tests/transformers/test_matrixgame2.py

@@ -1,133 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import os
-
-import pytest
-import torch
-
-from fastvideo.configs.models.dits.matrixgame2 import MatrixGame2WanVideoConfig
-from fastvideo.configs.pipelines import PipelineConfig
-from fastvideo.distributed.parallel_state import (
-    get_sp_parallel_rank,
-    get_sp_world_size,
-)
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.forward_context import set_forward_context
-from fastvideo.logger import init_logger
-from fastvideo.models.loader.component_loader import TransformerLoader
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.utils import maybe_download_model
-
-logger = init_logger(__name__)
-
-os.environ["MASTER_ADDR"] = "localhost"
-os.environ["MASTER_PORT"] = "29505"
-os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")
-
-MODEL_PATH = maybe_download_model("FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers")
-TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
-REFERENCE_LATENT = -111835.58895772696
-
-
-@pytest.mark.skip(reason="Not reliably reproducible")
-@pytest.mark.usefixtures("distributed_setup")
-def test_matrixgame2_transformer():
-    transformer_path = TRANSFORMER_PATH
-
-    sp_rank = get_sp_parallel_rank()
-    sp_world_size = get_sp_world_size()
-
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    precision = torch.bfloat16
-    precision_str = "bf16"
-
-    args = FastVideoArgs(
-        model_path=transformer_path,
-        dit_cpu_offload=False,
-        use_fsdp_inference=False,
-        pipeline_config=PipelineConfig(
-            dit_config=MatrixGame2WanVideoConfig(), dit_precision=precision_str
-        ),
-    )
-    args.device = device
-
-    loader = TransformerLoader()
-    model = loader.load(transformer_path, args).to(device, dtype=precision)
-    model.eval()
-
-    total_params = sum(p.numel() for p in model.parameters())
-    weight_sum = sum(p.to(torch.float64).sum().item() for p in model.parameters())
-    weight_mean = weight_sum / total_params
-    logger.info("Total parameters: %s", total_params)
-    logger.info("Weight sum: %s", weight_sum)
-    logger.info("Weight mean: %s", weight_mean)
-
-    torch.manual_seed(42)
-
-    batch_size = 1
-    latent_frames = 15
-    latent_height = 44
-    latent_width = 80
-    image_dim = 1280
-    image_seq_len = 257
-
-    # Set num_frame_per_block to match actual latent_frames for action_model
-    model.num_frame_per_block = latent_frames
-    # Force full attention mode
-    model.local_attn_size = -1
-    # Clear cached block masks so they regenerate with correct settings
-    model.block_mask = None
-    model.block_mask_mouse = None
-    model.block_mask_keyboard = None
-
-    # Video latents [B, C, T, H, W]
-    x = torch.randn(batch_size, 16, latent_frames, latent_height, latent_width,
-                    device=device, dtype=precision)
-    cond_concat = torch.randn(batch_size, 20, latent_frames, latent_height, latent_width,
-                              device=device, dtype=precision)
-    hidden_states = torch.cat([x, cond_concat], dim=1)
-
-    if sp_world_size > 1:
-        chunk_per_rank = hidden_states.shape[2] // sp_world_size
-        hidden_states = hidden_states[:, :, sp_rank * chunk_per_rank:(sp_rank + 1) * chunk_per_rank]
-
-    # Image embeddings
-    encoder_hidden_states_image = [
-        torch.randn(batch_size, image_seq_len, image_dim, device=device, dtype=precision)
-    ]
-
-    # Timestep
-    timestep = torch.full((batch_size, latent_frames), 500, device=device, dtype=precision)
-
-    # Action conditions
-    num_pixel_frames = (latent_frames - 1) * 4 + 1
-    mouse_cond = torch.randn(batch_size, num_pixel_frames, 2, device=device, dtype=precision)
-    keyboard_cond = torch.randn(batch_size, num_pixel_frames, 4, device=device, dtype=precision)
-
-    # Text embeddings (empty for i2v)
-    encoder_hidden_states = torch.zeros(batch_size, 0, model.config.hidden_size,
-                                        device=device, dtype=precision)
-
-    forward_batch = ForwardBatch(data_type="dummy")
-    model.num_frame_per_block = latent_frames
-
-    with torch.no_grad():
-        with torch.amp.autocast('cuda', dtype=precision):
-            with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=forward_batch):
-                output = model(
-                    hidden_states=hidden_states,
-                    encoder_hidden_states=encoder_hidden_states,
-                    encoder_hidden_states_image=encoder_hidden_states_image,
-                    timestep=timestep,
-                    mouse_cond=mouse_cond,
-                    keyboard_cond=keyboard_cond,
-                )
-
-    latent = output.double().sum().item()
-
-    diff = abs(REFERENCE_LATENT - latent)
-    relative_diff = diff / abs(REFERENCE_LATENT)
-    logger.info(f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}")
-    logger.info(f"Absolute diff: {diff}, Relative diff: {relative_diff * 100:.4f}%")
-
-    # The 0.19% diff now
-    assert relative_diff < 0.005, f"Output latents differ significantly: relative diff = {relative_diff * 100:.4f}% (max allowed: 0.5%)"

fastvideo/tests/vaes/test_hunyuan_vae.py

@@ -1,82 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-import json
-import os
-
-import pytest
-import torch
-from safetensors.torch import load_file
-
-from fastvideo.configs.pipelines import PipelineConfig
-from fastvideo.logger import init_logger
-
-# from fastvideo.models.vaes.hunyuanvae import (
-#     AutoencoderKLHunyuanVideo as MyHunyuanVAE)
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.models.loader.component_loader import VAELoader
-from fastvideo.configs.models.vaes import HunyuanVAEConfig
-from fastvideo.utils import maybe_download_model
-
-logger = init_logger(__name__)
-
-os.environ["MASTER_ADDR"] = "localhost"
-os.environ["MASTER_PORT"] = "29503"
-
-BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH))
-VAE_PATH = os.path.join(MODEL_PATH, "vae")
-CONFIG_PATH = os.path.join(VAE_PATH, "config.json")
-
-# Latent generated on commit d71a4ebffc2034922fc379568b6a6aa722f3744c with 1 x A40
-# torch 2.7.1
-A40_REFERENCE_LATENT = -106.22467041015625
-
-# Latent generated on commit 2b54068960c41d42221e8b8719a374b499855029 with 1 x L40S
-L40S_REFERENCE_LATENT = -157.09130859375
-
-
-@pytest.mark.usefixtures("distributed_setup")
-@pytest.mark.skip(reason="Skipping test_hunyuan_vae")
-def test_hunyuan_vae():
-    device = torch.device("cuda:0")
-    precision = torch.bfloat16
-    precision_str = "bf16"
-    args = FastVideoArgs(
-        model_path=VAE_PATH, pipeline_config=PipelineConfig(vae_config=HunyuanVAEConfig(), vae_precision=precision_str)
-    )
-    args.device = device
-    args.vae_cpu_offload = False
-
-    loader = VAELoader()
-    model = loader.load(VAE_PATH, args)
-
-    model.enable_tiling(
-        tile_sample_min_height=32,
-        tile_sample_min_width=32,
-        tile_sample_min_num_frames=8,
-        tile_sample_stride_height=16,
-        tile_sample_stride_width=16,
-        tile_sample_stride_num_frames=4,
-    )
-
-    batch_size = 1
-
-    # Video input [B, C, T, H, W]
-    input_tensor = torch.randn(batch_size, 3, 21, 64, 64, device=device, dtype=torch.bfloat16)
-
-    # Disable gradients for inference
-    with torch.no_grad():
-        latent = model.encode(input_tensor).mean.double().sum().item()
-
-    # Check if latents are similar
-    device_name = torch.cuda.get_device_name()
-    if "A40" in device_name:
-        REFERENCE_LATENT = A40_REFERENCE_LATENT
-    elif "L40S" in device_name:
-        REFERENCE_LATENT = L40S_REFERENCE_LATENT
-    else:
-        raise ValueError(f"Unknown device: {device_name}")
-
-    diff_encoded_latents = abs(REFERENCE_LATENT - latent)
-    logger.info(f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}")
-    assert diff_encoded_latents < 1e-4, f"Encoded latents differ significantly: max diff = {diff_encoded_latents}"

fastvideo/tests/vaes/test_official_wan_vae.py

@@ -1,113 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import os
-
-import numpy as np
-import pytest
-import torch
-from diffusers import AutoencoderKLWan
-
-from fastvideo.configs.pipelines import PipelineConfig
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.logger import init_logger
-from fastvideo.models.loader.component_loader import VAELoader
-from fastvideo.configs.models.vaes import WanVAEConfig
-from fastvideo.utils import maybe_download_model
-
-try:
-    from wan.modules.vae2_2 import Wan2_2_VAE
-except ImportError:
-    Wan2_2_VAE = None
-
-logger = init_logger(__name__)
-
-os.environ["MASTER_ADDR"] = "localhost"
-os.environ["MASTER_PORT"] = "29503"
-
-BASE_MODEL_PATH = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH)
-VAE_PATH = os.path.join(MODEL_PATH, "vae")
-
-
-@pytest.mark.skip(reason="disable test")
-@pytest.mark.usefixtures("distributed_setup")
-def test_official_wan_vae():
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    precision = torch.bfloat16
-    precision_str = "bf16"
-    args = FastVideoArgs(model_path=VAE_PATH, pipeline_config=PipelineConfig(vae_config=WanVAEConfig(), vae_precision=precision_str))
-    args.device = device
-    args.vae_cpu_offload = False
-
-    loader = VAELoader()
-    model2 = loader.load(VAE_PATH, args)
-    assert model2.use_feature_cache # Default to use the original WanVAE algorithm
-
-    model1 = Wan2_2_VAE(vae_pth="/mnt/weka/home/hao.zhang/wei/Wan2.2_VAE.pth", device=device, dtype=precision)
-
-    # Create identical inputs for both models
-    batch_size = 1
-
-    # Video input [B, C, T, H, W]
-    input_tensor = torch.randn(batch_size,
-                               3,
-                               81,
-                               32,
-                               32,
-                               device=device,
-                               dtype=precision)
-    # latent_tensor = torch.randn(batch_size,
-    #                             16,
-    #                             21,
-    #                             32,
-    #                             32,
-    #                             device=device,
-    #                             dtype=precision)
-
-    # Disable gradients for inference
-    with torch.no_grad():
-        # Test encoding
-        logger.info("Testing encoding...")
-        latent1 = model1.encode([input_tensor.squeeze(0)])[0].unsqueeze(0)
-        print("--------------------------------")
-        latent2 = model2.encode(input_tensor)
-        # Check if latents have the same shape
-        assert latent1.shape == latent2.mean.shape, f"Latent shapes don't match: {latent1.mean.shape} vs {latent2.mean.shape}"
-        # Check if latents are similar
-        max_diff_encode = torch.max(torch.abs(latent1 - latent2.mean))
-        mean_diff_encode = torch.mean(torch.abs(latent1 - latent2.mean))
-        logger.info("Maximum difference between encoded latents: %s",
-                    max_diff_encode.item())
-        logger.info("Mean difference between encoded latents: %s",
-                    mean_diff_encode.item())
-        assert max_diff_encode < 1e-5, f"Encoded latents differ significantly: max diff = {mean_diff_encode.item()}"
-        # Test decoding
-        logger.info("Testing decoding...")
-        # latent1_tensor = latent1.mode()
-        # mean1 = (torch.tensor(model1.config.latents_mean).view(
-        #     1, model1.config.z_dim, 1, 1, 1).to(input_tensor.device,
-        #                                         input_tensor.dtype))
-        # std1 = (1.0 / torch.tensor(model1.config.latents_std).view(
-        #     1, model1.config.z_dim, 1, 1, 1)).to(input_tensor.device,
-        #                                         input_tensor.dtype)
-        # latent1_tensor = latent1_tensor / std1 + mean1
-        # output1 = model1.decode(latent1_tensor).sample
-
-        output1 = model1.decode([latent1.squeeze(0)])
-        output1 = output1[0].unsqueeze(0)
-
-        mean2 = model2.config.arch_config.shift_factor.to(input_tensor.device, input_tensor.dtype)
-        std2 = model2.config.arch_config.scaling_factor.to(input_tensor.device, input_tensor.dtype)
-        latent2_tensor = latent2.mode()
-        latent2_tensor = latent2_tensor / std2 + mean2
-        output2 = model2.decode(latent2_tensor)
-        # Check if outputs have the same shape
-        assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}"
-
-        # Check if outputs are similar
-        max_diff_decode = torch.max(torch.abs(output1.float() - output2.float()))
-        mean_diff_decode = torch.mean(torch.abs(output1.float() - output2.float()))
-        logger.info("Maximum difference between decoded outputs: %s",
-                    max_diff_decode.item())
-        logger.info("Mean difference between decoded outputs: %s",
-                    mean_diff_decode.item())
-        assert max_diff_decode < 1e-5, f"Decoded outputs differ significantly: max diff = {max_diff_decode.item()}"

fastvideo/train/callbacks/validation.py

@@ -404,6 +404,7 @@ class ValidationCallback(Callback):
                         art = self.tracker.video(
                             fname,
                             caption=cap,
+                            fps=sp.fps,
                         )
                         if art is not None:
                             video_logs.append(art)

fastvideo/training/distillation_pipeline.py

@@ -1180,7 +1180,7 @@ class DistillationPipeline(TrainingPipeline):
 
                     artifacts = []
                     for filename, caption in zip(video_filenames, all_captions, strict=True):
-                        video_artifact = self.tracker.video(filename, caption=caption)
+                        video_artifact = self.tracker.video(filename, caption=caption, fps=sampling_param.fps)
                         if video_artifact is not None:
                             artifacts.append(video_artifact)
                     if artifacts:

fastvideo/training/trackers.py

@@ -7,22 +7,28 @@ interface that can be used across all FastVideo training pipelines.
 
 from __future__ import annotations
 
+from collections.abc import Iterable, Iterator
 import contextlib
 import copy
+import math
 import os
 import pathlib
+import tempfile
 import time
 from dataclasses import dataclass
 from enum import Enum
 from typing import Any
-from collections.abc import Iterable, Iterator
 
+import numpy as np
 import torch
 
 from fastvideo.logger import init_logger
 
 logger = init_logger(__name__)
 
+_DEFAULT_VIDEO_FPS = 16
+_MISSING_ARTIFACT = object()
+
 
 def _sanitize_wandb_config(value: Any) -> Any:
     """Best-effort conversion of nested config objects to W&B-safe values."""
@@ -56,6 +62,140 @@ def _sanitize_wandb_config(value: Any) -> Any:
     return repr(value)
 
 
+def _prepare_video_array(data: Any) -> np.ndarray:
+    """Convert W&B-style TCHW/BTCHW video data into GIF-ready frames."""
+    if isinstance(data, torch.Tensor):
+        data = data.detach().cpu().numpy()
+
+    video = np.asarray(data)
+    if video.ndim == 4:
+        video = video.reshape(1, *video.shape)
+    elif video.ndim != 5:
+        raise ValueError("Video data must have shape [T, C, H, W] or [B, T, C, H, W]")
+
+    batch_size, num_frames, channels, height, width = video.shape
+    if batch_size == 0 or num_frames == 0:
+        raise ValueError("Video data must contain at least one batch item and one frame")
+    if channels not in (1, 3, 4):
+        raise ValueError(f"Video data must have 1, 3, or 4 channels; got {channels}")
+    if video.dtype != np.uint8:
+        logger.warning("Converting video data to uint8 for SwanLab")
+        video = video.astype(np.uint8)
+
+    # Match wandb.Video's batch tiling so the same input has a familiar layout
+    # in either tracker.
+    if batch_size & (batch_size - 1):
+        padded_batch_size = 1 << batch_size.bit_length()
+        padding = np.zeros(
+            (padded_batch_size - batch_size, num_frames, channels, height, width),
+            dtype=video.dtype,
+        )
+        video = np.concatenate((video, padding), axis=0)
+
+    num_rows = 1 << ((batch_size.bit_length() - 1) // 2)
+    num_columns = video.shape[0] // num_rows
+    video = video.reshape(num_rows, num_columns, num_frames, channels, height, width)
+    video = np.transpose(video, axes=(2, 0, 4, 1, 5, 3))
+    video = video.reshape(num_frames, num_rows * height, num_columns * width, channels)
+    if channels == 1:
+        video = video[..., 0]
+    return np.ascontiguousarray(video)
+
+
+def _read_video_file(file_path: str) -> tuple[list[np.ndarray], float | None]:
+    """Decode a video file and return its frames and encoded frame rate."""
+    import imageio.v2 as imageio
+
+    reader = imageio.get_reader(file_path)
+    try:
+        metadata = reader.get_meta_data() or {}
+        frames = [np.asarray(frame) for frame in reader]
+    finally:
+        reader.close()
+
+    if not frames:
+        raise ValueError(f"Video file contains no frames: {file_path}")
+
+    raw_source_fps = metadata.get("fps")
+    try:
+        source_fps = float(str(raw_source_fps))
+    except ValueError:
+        source_fps = None
+    if source_fps is not None and (not math.isfinite(source_fps) or source_fps <= 0):
+        source_fps = None
+    return frames, source_fps
+
+
+def _coerce_video_fps(fps: int | float | None) -> float:
+    if fps is None:
+        return _DEFAULT_VIDEO_FPS
+    value = float(fps)
+    if not math.isfinite(value) or value <= 0:
+        raise ValueError(f"Video fps must be a positive finite number; got {fps!r}")
+    return value
+
+
+def _write_gif(file_path: str, frames: Any, fps: float) -> None:
+    """Encode frames as an animated GIF for SwanLab's video API."""
+    from PIL import Image
+
+    images = []
+    for frame in frames:
+        array = np.asarray(frame)
+        if array.dtype != np.uint8:
+            array = array.astype(np.uint8)
+        if array.ndim == 3 and array.shape[-1] == 1:
+            array = array[..., 0]
+        if array.ndim not in (2, 3) or (array.ndim == 3 and array.shape[-1] not in (3, 4)):
+            raise ValueError(f"GIF frames must be grayscale, RGB, or RGBA; got shape {array.shape}")
+        images.append(Image.fromarray(array))
+
+    if not images:
+        raise ValueError("Video data must contain at least one frame")
+
+    duration_ms = max(1, round(1000 / fps))
+    images[0].save(
+        file_path,
+        format="GIF",
+        save_all=True,
+        append_images=images[1:],
+        duration=duration_ms,
+        loop=0,
+    )
+
+
+@dataclass(frozen=True)
+class _SequentialArtifact:
+    """Backend-specific versions of one artifact created by a tracker group."""
+
+    values: tuple[Any | None, ...]
+
+
+def _select_tracker_artifact(value: Any, tracker_index: int) -> Any:
+    """Resolve nested sequential artifacts for one child tracker."""
+    if isinstance(value, _SequentialArtifact):
+        selected = value.values[tracker_index]
+        return _MISSING_ARTIFACT if selected is None else selected
+    if isinstance(value, dict):
+        selected_dict = {}
+        for key, item in value.items():
+            selected = _select_tracker_artifact(item, tracker_index)
+            if selected is not _MISSING_ARTIFACT:
+                selected_dict[key] = selected
+        return selected_dict if selected_dict else _MISSING_ARTIFACT
+    if isinstance(value, list):
+        selected_list = [
+            selected for item in value
+            if (selected := _select_tracker_artifact(item, tracker_index)) is not _MISSING_ARTIFACT
+        ]
+        return selected_list if selected_list else _MISSING_ARTIFACT
+    if isinstance(value, tuple):
+        selected_tuple = tuple(selected for item in value
+                               if (selected := _select_tracker_artifact(item, tracker_index)) is not _MISSING_ARTIFACT)
+        return selected_tuple if selected_tuple else _MISSING_ARTIFACT
+    return value
+
+
 @dataclass
 class Timer:
     """Simple timer utility used by the trackers."""
@@ -222,8 +362,7 @@ class WandbTracker(BaseTracker):
         kwargs: dict[str, Any] = {}
         if caption is not None:
             kwargs["caption"] = caption
-        if fps is not None:
-            kwargs["fps"] = fps
+        kwargs["fps"] = fps if fps is not None else _DEFAULT_VIDEO_FPS
         if format is not None:
             kwargs["format"] = format
         else:
@@ -254,8 +393,10 @@ class SequentialTracker(BaseTracker):
         self._timed_metrics = {}
 
     def log_artifacts(self, artifacts: dict[str, Any], step: int) -> None:
-        for tracker in self._trackers:
-            tracker.log_artifacts(artifacts, step)
+        for tracker_index, tracker in enumerate(self._trackers):
+            tracker_artifacts = _select_tracker_artifact(artifacts, tracker_index)
+            if tracker_artifacts is not _MISSING_ARTIFACT:
+                tracker.log_artifacts(tracker_artifacts, step)
         self._timed_metrics = {}
 
     def log_file(
@@ -278,16 +419,91 @@ class SequentialTracker(BaseTracker):
         fps: int | None = None,
         format: str | None = None,
     ) -> Any | None:
-        for tracker in self._trackers:
-            video = tracker.video(data, caption=caption, fps=fps, format=format)
-            if video is not None:
-                return video
-        return None
+        videos = tuple(tracker.video(data, caption=caption, fps=fps, format=format) for tracker in self._trackers)
+        if all(video is None for video in videos):
+            return None
+        return _SequentialArtifact(videos)
+
+
+class SwanlabTracker(BaseTracker):
+    """Tracker implementation for SwanLab."""
+
+    def __init__(
+        self,
+        experiment_name: str,
+        log_dir: str,
+        *,
+        config: dict[str, Any] | None = None,
+        run_name: str | None = None,
+    ) -> None:
+        super().__init__()
+
+        try:
+            import swanlab
+        except ModuleNotFoundError as error:
+            if error.name != "swanlab":
+                raise
+            raise ModuleNotFoundError("SwanLab tracking requires the optional 'swanlab' dependency. "
+                                      "Install it with `uv pip install 'fastvideo[swanlab]'` (or "
+                                      "`uv pip install -e '.[swanlab]'` from a source checkout).") from error
+
+        pathlib.Path(log_dir).mkdir(parents=True, exist_ok=True)
+
+        self._swanlab = swanlab
+        self._run = swanlab.init(
+            project=experiment_name,
+            experiment_name=run_name,
+            config=(_sanitize_wandb_config(config) if config is not None else None),
+            logdir=log_dir,
+        )
+        logger.info("Initialized SwanLab tracker")
+
+    def log(self, metrics: dict[str, Any], step: int) -> None:
+        metrics = {**self._timed_metrics, **metrics}
+        if metrics:
+            self._swanlab.log(metrics, step=step)
+        self._timed_metrics = {}
+
+    def finish(self) -> None:
+        self._swanlab.finish()
+
+    def video(
+        self,
+        data: Any,
+        *,
+        caption: str | None = None,
+        fps: int | None = None,
+        format: str | None = None,
+    ) -> Any:
+        """Create a SwanLab GIF artifact from a file or W&B-style array."""
+        del format  # SwanLab currently supports GIF artifacts only.
+
+        if isinstance(data, str | os.PathLike):
+            source_path = os.fspath(data)
+            if pathlib.Path(source_path).suffix.lower() == ".gif":
+                return self._swanlab.Video(source_path, caption=caption)
+            frames, source_fps = _read_video_file(source_path)
+            if fps is not None:
+                video_fps = _coerce_video_fps(fps)
+            else:
+                video_fps = source_fps if source_fps is not None else _coerce_video_fps(None)
+        else:
+            frames = _prepare_video_array(data)
+            video_fps = _coerce_video_fps(fps)
+
+        file_descriptor, gif_path = tempfile.mkstemp(suffix=".gif")
+        os.close(file_descriptor)
+        try:
+            _write_gif(gif_path, frames, video_fps)
+            return self._swanlab.Video(gif_path, caption=caption)
+        finally:
+            pathlib.Path(gif_path).unlink(missing_ok=True)
 
 
 class Trackers(str, Enum):
     NONE = "none"
     WANDB = "wandb"
+    SWANLAB = "swanlab"
 
 
 SUPPORTED_TRACKERS = {tracker.value for tracker in Trackers}
@@ -324,6 +540,14 @@ def initialize_trackers(
                     config=config,
                     run_name=run_name,
                 ))
+        elif tracker_name == Trackers.SWANLAB.value:
+            tracker_instances.append(
+                SwanlabTracker(
+                    experiment_name,
+                    os.path.abspath(log_dir),
+                    config=config,
+                    run_name=run_name,
+                ))
 
     if not tracker_instances:
         return DummyTracker()

fastvideo/training/training_pipeline.py

@@ -831,7 +831,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
 
                 artifacts = []
                 for filename, caption in zip(video_filenames, all_captions, strict=True):
-                    video_artifact = self.tracker.video(filename, caption=caption)
+                    video_artifact = self.tracker.video(filename, caption=caption, fps=sampling_param.fps)
                     if video_artifact is not None:
                         artifacts.append(video_artifact)
                 if artifacts:

mkdocs.yml

@@ -169,6 +169,7 @@ nav:
   - Training:
     - Overview: training/overview.md
     - Training Infrastructure: training/train_infra.md
+    - Experiment Tracking: training/trackers.md
     - Data Preprocessing: training/data_preprocess.md
     - Fine-tuning: training/finetune.md
     - Examples:

pyproject.toml

@@ -114,12 +114,15 @@ fastvideo-kernel = [
 imagebind = { git = "https://github.com/facebookresearch/ImageBind.git", rev = "53680b02d7e37b19b124fa37bae4b6c98c38f5be" }
 # FA4 cute, pinned to a cutlass-4.5-compatible revision. torch.compile support
 # comes from FastVideo's own custom_op wrappers.
-flash-attn-4 = { git = "https://github.com/Dao-AILab/flash-attention.git", rev = "940cd9680f3315f2f06b43ab5bea2c2cf2d96806", subdirectory = "flash_attn/cute" }
+flash-attn-4 = { git = "https://github.com/Dao-AILab/flash-attention.git", rev = "82d6441eec5d4dfec120153db2c0145ae855a083", subdirectory = "flash_attn/cute" }
 
 [project.optional-dependencies]
 
 # flash-attn: uv pip install flash-attn==2.8.1 --no-cache-dir --no-build-isolation
 
+swanlab = [
+    "swanlab>=0.6.7",
+]
 
 lint = [
     "pre-commit==4.0.1",

tests/local_tests/test_nvfp4_fa4.py

@@ -10,15 +10,6 @@ import pytest
 
 os.environ["CUTE_DSL_ENABLE_TVM_FFI"] = "1"
 
-# Warm up flashinfer JIT before any FA4 imports
-from flashinfer.quantization import nvfp4_quantize, SfLayout
-_warmup = nvfp4_quantize(
-    torch.randn(4, 128, device="cuda", dtype=torch.bfloat16),
-    torch.ones(1, device="cuda", dtype=torch.float32),
-    sfLayout=SfLayout.layout_128x4, do_shuffle=False,
-)
-del _warmup
-
 
 def _make_impl(nheads, headdim, nvfp4=False):
     from fastvideo.attention.backends.flash_attn import FlashAttentionImpl
@@ -45,7 +36,8 @@ def _cuda_timer(fn, warmup=5, iters=20):
 
 
 @pytest.mark.skipif(
-    torch.cuda.get_device_capability() not in [(10, 0), (10, 3)],
+    not torch.cuda.is_available()
+    or torch.cuda.get_device_capability() not in [(10, 0), (10, 3)],
     reason="Requires Blackwell GPU (sm100a or sm103a)"
 )
 class TestNVFP4FA4:
@@ -55,6 +47,17 @@ class TestNVFP4FA4:
     MODEL_HEADDIM = 128
     MODEL_SEQLEN = 32760  # 480x832 video, 81 frames
 
+    @pytest.fixture(scope="class", autouse=True)
+    def _flashinfer_warmup(self):
+        # Warm up flashinfer JIT before any FA4 imports
+        from flashinfer.quantization import SfLayout, nvfp4_quantize
+        _warmup = nvfp4_quantize(
+            torch.randn(4, 128, device="cuda", dtype=torch.bfloat16),
+            torch.ones(1, device="cuda", dtype=torch.float32),
+            sfLayout=SfLayout.layout_128x4, do_shuffle=False,
+        )
+        del _warmup
+
     def test_quantize_fn_shapes(self):
         """_nvfp4_quantize_for_fa4 produces correct shapes and strides."""
         from fastvideo.attention.backends.flash_attn import _nvfp4_quantize_for_fa4