Compare commits

...
Author SHA1 Message Date
SolitaryThinker bb3b3d0657 [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.
2026-07-06 23:10:15 -07:00
SolitaryThinker 67a8f7ac85 [fix]: complete SwanLab tracker integration 2026-07-06 23:10:15 -07:00
kared d133aef55d [feat]: add swanlab tracker 2026-07-06 23:10:15 -07:00
William Lin d921832cd2 [misc]: reserve CPU/memory for timing-sensitive Modal CI lanes (#1566) 2026-07-06 22:13:53 -07:00
William Lin 629697629a [bugfix]: free CUDA memory between train-framework model tests (LongCat OOM on L40S) (#1565) 2026-07-06 21:27:30 -07:00
William Lin a25313beec [ci]: wire fastvideo/tests/ops/ into the unit-test lane (#1559) 2026-07-06 12:07:26 -07:00
William Lin dbde64385b [bugfix]: bump FA4 pin to the CuTe DSL 4.6 compatible rev (#1564) 2026-07-06 12:06:59 -07:00
William Lin 9d909f5f04 [test]: remove dead and duplicate tests (-489 lines) (#1556) 2026-07-05 15:53:40 -07:00
William Lin 76b0550c15 [ci]: run pre-commit on fork PRs without manual approval (#1555) 2026-07-05 14:18:16 -07:00
William Lin 384c1e9493 [misc]: update reseed-performance-baseline skill for the hf_store move (#1545 follow-up) (#1553) 2026-07-05 14:16:55 -07:00
William Lin b1dbcc93f6 [misc]: reformat fastvideo/performance to the repo yapf config (#1554) 2026-07-05 14:16:20 -07:00
William Lin b93833772e [ci]: guard against test directories no CI lane collects (#1552) 2026-07-05 14:07:38 -07:00
Mac Lee 30b523edd6 [ci] Normalize performance stage component metrics (#1475) (#1550) 2026-07-05 14:05:26 -07:00
31 changed files with 780 additions and 549 deletions
@@ -69,7 +69,7 @@ approval, then upload reviewed accepted baseline records.
| `model_id` | Yes | Benchmark id, e.g. `wan-t2v-1.3b-2gpu`. This maps to the HF subdirectory after `sanitize(model_id)`. |
| `gpu_type` | Yes | Exact GPU device string from the performance record, e.g. the L40S device name emitted by CI. Baselines are GPU-specific. |
| `source_results` | Yes | One or more local paths or Buildkite artifact URLs for accepted shifted performance JSONs. Prefer normalized `normalized_perf_*.json` artifacts emitted by `compare_baseline.py`. Accept `source_result` as an alias only for a single JSON. |
| `max_intra_batch_regression` | No | Maximum allowed regression of any source JSON against the source batch median. Default: `PERF_MAX_REGRESSION` if set, otherwise `0.05` (5%). |
| `max_intra_batch_regression` | No | Maximum allowed regression of any source JSON against the source batch median. Default: `0.05` (5%). |
| `intent_rationale` | Yes | One-line explanation for why the baseline shift is legitimate. This is written into provenance and should be reused in the PR. |
Hardcoded defaults:
@@ -148,8 +148,7 @@ For each metric with at least two non-null source values:
4. Stop if any source record regresses against the batch median by more than
`max_intra_batch_regression`.
Default `max_intra_batch_regression` to `PERF_MAX_REGRESSION` when set,
otherwise `0.05`. Print a table with per-source values, batch median, and
Default `max_intra_batch_regression` to `0.05`. Print a table with per-source values, batch median, and
worst intra-batch regression.
This check prevents uploading a mixed batch where one JSON is materially
@@ -183,7 +182,7 @@ present, that run is not a valid source for baseline reseeding.
### 2. Sync and back up existing HF records under /tmp
Use `fastvideo/tests/performance/hf_store.py` helpers directly. Do **not** use
Use `fastvideo/performance/hf_store.py` helpers directly. Do **not** use
`compare_baseline.py` as a sync shortcut; on full main runs it can persist
records, while this step must only fetch and back up existing history.
@@ -192,7 +191,7 @@ The sync command pattern is:
```bash
export PERFORMANCE_TRACKING_ROOT="${PERFORMANCE_TRACKING_ROOT:-/tmp/perf-tracking}"
export HF_REPO_ID="${HF_REPO_ID:-FastVideo/performance-tracking}"
PYTHONPATH=fastvideo/tests/performance python -c 'from hf_store import sync_from_hf; import os; sync_from_hf(os.environ["PERFORMANCE_TRACKING_ROOT"], strict=True)'
python -c 'from fastvideo.performance.hf_store import sync_from_hf; import os; sync_from_hf(os.environ["PERFORMANCE_TRACKING_ROOT"], strict=True)'
```
Then back up only the sanitized model directory under `/tmp`:
@@ -200,8 +199,8 @@ Then back up only the sanitized model directory under `/tmp`:
```bash
SHORT_COMMIT=$(git rev-parse --short=12 HEAD)
TIMESTAMP=$(date -u +%Y%m%d_%H%M%S)
MODEL_SAFE=$(PYTHONPATH=fastvideo/tests/performance python - <<'PY'
from hf_store import sanitize
MODEL_SAFE=$(python - <<'PY'
from fastvideo.performance.hf_store import sanitize
print(sanitize("<model_id>"))
PY
)
@@ -235,7 +234,7 @@ first baseline seed. Continue, but report that baseline history was empty.
Load the last 5 successful records for the target:
```python
from hf_store import load_records_for_model
from fastvideo.performance.hf_store import load_records_for_model
records = load_records_for_model(
"/tmp/perf-tracking",
@@ -372,7 +371,7 @@ prepared records plus backup on disk.
Use the shared storage helper so the path and repo type match CI:
```python
from hf_store import upload_record
from fastvideo.performance.hf_store import upload_record
upload_record("<local_record_path>", record, strict=True)
```
@@ -460,7 +459,7 @@ directories created for this reseed. Never remove unrelated `/tmp` contents.
intentional baseline replacement.
- `fastvideo/tests/performance/compare_baseline.py` — normalization, rolling
median comparison, and persistence rules.
- `fastvideo/tests/performance/hf_store.py` — HF sync, record loading,
- `fastvideo/performance/hf_store.py` — HF sync, record loading,
`sanitize()`, and `upload_record()`.
- `fastvideo/tests/performance/test_inference_performance.py` — source result
JSON schema.
+19 -2
View File
@@ -1,7 +1,11 @@
name: pre-commit
on:
pull_request:
# pull_request_target instead of pull_request: the workflow definition and
# the hook config are always taken from the BASE branch, so fork /
# first-time-contributor PRs run immediately without a maintainer clicking
# "Approve and run". The PR head is checked out as data only.
pull_request_target:
branches: [main]
workflow_call:
inputs:
@@ -15,12 +19,25 @@ permissions:
jobs:
pre-commit:
if: github.event_name == 'workflow_call' || github.event.pull_request.draft != true
if: github.event.pull_request.draft != true
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
with:
ref: ${{ inputs.ref || '' }}
# For PR events, lint the PR head — but keep the hook definitions from
# the base branch so an untrusted PR cannot alter what gets executed.
- name: Save trusted hook config
if: github.event_name == 'pull_request_target'
run: cp .pre-commit-config.yaml "$RUNNER_TEMP/trusted-pre-commit-config.yaml"
- uses: actions/checkout@v4
if: github.event_name == 'pull_request_target'
with:
ref: ${{ github.event.pull_request.head.sha }}
persist-credentials: false
- name: Restore trusted hook config
if: github.event_name == 'pull_request_target'
run: cp "$RUNNER_TEMP/trusted-pre-commit-config.yaml" .pre-commit-config.yaml
- uses: actions/setup-python@v5
with:
python-version: "3.12"
+2 -2
View File
@@ -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
+7 -6
View File
@@ -112,10 +112,11 @@ per-metric policy with direction, percent threshold, absolute threshold, and a
while it runs so pipeline stage execution times are available in
`generate_video(...).logging_info`. Stage logs use pipeline-unique keys such as
`prompt_encoding_stage` so duplicate stage classes do not collide. For
`PipelineStage` entries, the extractor maps the `stage_class` field:
`TextEncodingStage` maps to `text_encoder_time_s`, `DenoisingStage` and
`DmdDenoisingStage` map to `dit_time_s`, and `DecodingStage` maps to
`vae_decode_time_s`, with a fallback for older logs that used the class name as
`PipelineStage` entries, shared component stage bases emit a stable
`component_metric`: text encoding stages map to `text_encoder_time_s`,
denoising stages and subclasses map to `dit_time_s`, and decoding stages map to
`vae_decode_time_s`. The extractor falls back to known `stage_class` names for
older logs that do not include `component_metric` or that used the class name as
the stage key. Generator-side timings such as `PostDecodeFrameProcessStage`,
`VideoSaveStage`, and `AudioMuxStage` are intentionally ignored. If a pipeline
does not report one of the mapped stages, that component metric is stored as
@@ -432,5 +433,5 @@ pipelines that did not report a mapped component stage.
**Component timing is `null`** — the generated result did not include a mapped
stage in `logging_info.stages`. Check that the pipeline emits stage logging
and that the stage name is listed in `STAGE_METRIC_MAP` in
`test_inference_performance.py`.
and that the stage emits `component_metric` or is covered by the legacy
`STAGE_METRIC_MAP` fallback in `test_inference_performance.py`.
+91
View File
@@ -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).
+3
View File
@@ -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
+8
View File
@@ -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]:
+6 -2
View File
@@ -310,8 +310,12 @@ def load_records_for_model(
# ---------------------------------------------------------------------------
_NUMERIC_COLS = (
"latency", "throughput", "memory",
"text_encoder_time_s", "dit_time_s", "vae_decode_time_s",
"latency",
"throughput",
"memory",
"text_encoder_time_s",
"dit_time_s",
"vae_decode_time_s",
)
+8 -22
View File
@@ -36,6 +36,7 @@ DEFAULT_METRIC_POLICIES: tuple[MetricPolicy, ...] = (
MetricPolicy("vae_decode_time_s", "VAE Decode", 3, True, 0.05, 0.25),
)
def _optional_float(value: Any) -> float | None:
if value is None or isinstance(value, bool):
return None
@@ -57,9 +58,7 @@ def _optional_bool(value: Any) -> bool | None:
return None
def resolve_metric_policies(
threshold_overrides: Mapping[str, Any] | None,
) -> tuple[MetricPolicy, ...]:
def resolve_metric_policies(threshold_overrides: Mapping[str, Any] | None, ) -> tuple[MetricPolicy, ...]:
"""Return default metric policies with optional per-metric overrides."""
if not isinstance(threshold_overrides, Mapping):
@@ -80,25 +79,15 @@ def resolve_metric_policies(
label=base_policy.label,
precision=base_policy.precision,
lower_is_better=base_policy.lower_is_better,
threshold_percent=(
base_policy.threshold_percent
if threshold_percent is None
else threshold_percent
),
threshold_absolute=(
base_policy.threshold_absolute
if threshold_absolute is None
else threshold_absolute
),
threshold_percent=(base_policy.threshold_percent if threshold_percent is None else threshold_percent),
threshold_absolute=(base_policy.threshold_absolute
if threshold_absolute is None else threshold_absolute),
gated=base_policy.gated if gated is None else gated,
)
)
))
return tuple(policies)
def serialize_metric_thresholds(
policies: tuple[MetricPolicy, ...],
) -> dict[str, dict[str, float | bool]]:
def serialize_metric_thresholds(policies: tuple[MetricPolicy, ...], ) -> dict[str, dict[str, float | bool]]:
return {
policy.key: {
"threshold_percent": policy.threshold_percent,
@@ -118,10 +107,7 @@ def regression_delta(
return None
absolute_delta = current - baseline if policy.lower_is_better else baseline - current
percent_delta = absolute_delta / baseline
threshold_exceeded = (
percent_delta > policy.threshold_percent
and absolute_delta > policy.threshold_absolute
)
threshold_exceeded = (percent_delta > policy.threshold_percent and absolute_delta > policy.threshold_absolute)
return MetricDelta(
absolute=absolute_delta,
percent=percent_delta,
+4
View File
@@ -34,6 +34,7 @@ class PipelineStage(ABC):
composed with other stages to create a complete pipeline. Each stage is responsible
for a specific part of the process, such as prompt encoding, latent preparation, etc.
"""
performance_component_metric: str | None = None
def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
"""
@@ -155,6 +156,9 @@ class PipelineStage(ABC):
logger.info("[%s] Execution completed in %s ms", stage_name, execution_time * 1000)
batch.logging_info.add_stage_execution_time(stage_key, execution_time)
batch.logging_info.add_stage_metric(stage_key, "stage_class", stage_class_name)
component_metric = self.performance_component_metric
if component_metric is not None:
batch.logging_info.add_stage_metric(stage_key, "component_metric", component_metric)
except Exception as e:
torch.cuda.synchronize()
execution_time = time.perf_counter() - start_time
+1
View File
@@ -28,6 +28,7 @@ class DecodingStage(PipelineStage):
This stage handles the decoding of latent representations into the final
output format (e.g., pixel values).
"""
performance_component_metric = "vae_decode_time_s"
def __init__(self, vae, pipeline=None) -> None:
self.vae: ParallelTiledVAE = vae
+2
View File
@@ -51,6 +51,7 @@ class DenoisingStage(PipelineStage):
This stage handles the iterative denoising process that transforms
the initial noise into the final output.
"""
performance_component_metric = "dit_time_s"
def __init__(self, transformer, scheduler, pipeline=None, transformer_2=None, vae=None) -> None:
super().__init__()
@@ -1187,6 +1188,7 @@ class Cosmos25V2WDenoisingStage(Cosmos25DenoisingStage):
class Cosmos25AutoDenoisingStage(PipelineStage):
"""Route Cosmos 2.5 denoising to T2W vs V2W/I2W."""
performance_component_metric = "dit_time_s"
def __init__(self, transformer, scheduler) -> None:
super().__init__()
@@ -24,6 +24,7 @@ class TextEncodingStage(PipelineStage):
This stage handles the encoding of text prompts into the embedding space
expected by the diffusion model.
"""
performance_component_metric = "text_encoder_time_s"
def __init__(self, text_encoders, tokenizers) -> None:
"""
@@ -350,6 +351,7 @@ class Cosmos25TextEncodingStage(PipelineStage):
Cosmos 2.5 uses Reason1 (Qwen2.5-VL) and relies on the encoder's
`compute_text_embeddings_online()`.
"""
performance_component_metric = "text_encoder_time_s"
def __init__(self, text_encoder) -> None:
super().__init__()
@@ -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",
@@ -0,0 +1,84 @@
# SPDX-License-Identifier: Apache-2.0
"""Guard: every test directory must be collected by some CI lane or be on
the explicit allowlist below.
Three separate incidents on 2026-07-05 found test files that no CI lane
ever collects (fastvideo/tests/stages/, tests/local_tests/ additions in
PR #1509, and this sweep found seven dark directories in total): the tests
pass review, merge, and then silently never run. This test makes going
dark an explicit, reviewed decision instead of an accident: adding a new
test directory fails CI until it is either wired into a lane or
allowlisted here with a reason.
Pure text analysis — no fastvideo imports, no GPU, no torch.
"""
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parents[3]
TESTS_ROOT = REPO_ROOT / "fastvideo" / "tests"
# Files whose text constitutes "a CI lane references this directory".
CI_SOURCES = [
TESTS_ROOT / "modal" / "pr_test.py",
TESTS_ROOT / "modal" / "ssim_test.py",
*sorted((REPO_ROOT / ".buildkite").rglob("*.yml")),
*sorted((REPO_ROOT / ".buildkite").rglob("*.sh")),
]
# Directories that intentionally have no CI lane today. Every entry needs a
# reason; remove the entry when the directory gets wired into a lane.
# State as found on 2026-07-05 — these SHOULD shrink over time, not grow.
ALLOWLIST = {
"attention": "no lane yet — GPU attention-backend tests, run manually",
"audio": "no lane yet — audio encoder tests, run manually",
"distributed": "no lane yet — multi-GPU torchrun tests, run manually",
"hooks": "no lane yet — run manually",
"layers": "no lane yet — torchrun FSDP dispatch tests, run manually",
"nightly": "by design: nightly cadence, not per-PR",
"modal": "CI infrastructure itself, not a test suite",
}
def _dirs_with_tests() -> list[str]:
dirs = []
for child in sorted(TESTS_ROOT.iterdir()):
if child.is_dir() and any(child.rglob("test_*.py")):
dirs.append(child.name)
return dirs
def _ci_text() -> str:
return "\n".join(
src.read_text(errors="replace") for src in CI_SOURCES if src.exists())
def test_every_test_directory_is_collected_or_allowlisted():
ci_text = _ci_text()
dark = [
name for name in _dirs_with_tests()
if f"tests/{name}" not in ci_text and name not in ALLOWLIST
]
assert not dark, (
f"Test directories not referenced by any CI lane and not "
f"allowlisted: {dark}. Wire them into a lane in "
f"fastvideo/tests/modal/pr_test.py (or a Buildkite step), or add an "
f"allowlist entry with a reason in {__file__}.")
def test_local_tests_stays_out_of_ci():
# tests/local_tests/ (repo root) is developer-local by design (author
# decision, 2026-07-05): parity scaffolds and machine-specific checks
# that must never gate CI. Fail if any CI source starts collecting it.
assert "tests/local_tests" not in _ci_text(), (
"tests/local_tests/ is local-only by design; remove the CI "
"reference or move the tests into a fastvideo/tests/ lane.")
def test_allowlist_entries_are_still_real_directories():
# A stale allowlist hides regressions; entries must track reality.
missing = [
name for name in ALLOWLIST
if name != "modal" and not (TESTS_ROOT / name).is_dir()
]
assert not missing, (
f"Allowlisted directories no longer exist — remove them: {missing}")
+9 -1
View File
@@ -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=[
@@ -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):
@@ -179,7 +179,11 @@ def _extract_component_times(result: dict) -> dict[str, float | None]:
logger.debug("Skipping malformed stage '%s' data: %r", stage_name, stage_data)
continue
stage_class = stage_data.get("stage_class", stage_name)
metric_key = STAGE_METRIC_MAP.get(stage_class)
component_metric = stage_data.get("component_metric")
if isinstance(component_metric, str) and component_metric in component_times:
metric_key = component_metric
else:
metric_key = STAGE_METRIC_MAP.get(stage_class)
if metric_key is None:
logger.debug("Unmapped stage '%s' class '%s' (%.3fs)",
stage_name,
@@ -1,13 +1,20 @@
# SPDX-License-Identifier: Apache-2.0
from fastvideo.pipelines.pipeline_batch_info import PipelineLoggingInfo
from fastvideo.pipelines.stages.denoising import Cosmos25AutoDenoisingStage, DenoisingStage
from fastvideo.pipelines.stages.text_encoding import Cosmos25TextEncodingStage
from fastvideo.tests.performance.test_inference_performance import _extract_component_times
class SubclassStyleDenoisingStage(DenoisingStage):
pass
def test_extract_component_times_handles_pipeline_logging_info_object():
logging_info = PipelineLoggingInfo()
logging_info.add_stage_execution_time("prompt_encoding_stage", 1.25)
logging_info.add_stage_metric("prompt_encoding_stage", "stage_class", "TextEncodingStage")
logging_info.add_stage_metric("prompt_encoding_stage", "component_metric", "text_encoder_time_s")
assert _extract_component_times({"logging_info": logging_info}) == {
"text_encoder_time_s": 1.25,
@@ -45,6 +52,35 @@ def test_extract_component_times_uses_stage_class_for_pipeline_stage_keys():
}
def test_denoising_stage_subclasses_inherit_component_metric():
assert SubclassStyleDenoisingStage.performance_component_metric == "dit_time_s"
def test_cosmos25_direct_pipeline_stages_define_component_metrics():
assert Cosmos25TextEncodingStage.performance_component_metric == "text_encoder_time_s"
assert Cosmos25AutoDenoisingStage.performance_component_metric == "dit_time_s"
def test_extract_component_times_uses_component_metric_for_stage_subclasses():
result = {
"logging_info": {
"stages": {
"denoising_stage": {
"execution_time": 4.2,
"stage_class": "CosmosDenoisingStage",
"component_metric": "dit_time_s",
},
},
},
}
assert _extract_component_times(result) == {
"text_encoder_time_s": None,
"dit_time_s": 4.2,
"vae_decode_time_s": None,
}
def test_extract_component_times_keeps_legacy_class_name_keys():
# Backward-compatibility check for logs produced before pipeline-unique
# stage keys carried a separate stage_class field.
+225
View File
@@ -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"]}]
@@ -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}"
@@ -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%)"
-82
View File
@@ -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}"
@@ -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()}"
+1
View File
@@ -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)
+1 -1
View File
@@ -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:
+234 -10
View File
@@ -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()
+1 -1
View File
@@ -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:
+1
View File
@@ -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:
+4 -1
View File
@@ -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",
+13 -10
View File
@@ -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