nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix

Browse Source
This commit is contained in:
RandNMR73
2026-04-09 08:30:34 +00:00
parent eccf1a91d1
commit 10d99ea37e
13 changed files with 0 additions and 1301 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/tests/attention/test_flash_attn_cute_custom_op.py
-223
View File
@@ -1,223 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import importlib.util
import sys
from pathlib import Path
import pytest
import torch
_MODULE_NAME = "_fastvideo_flash_attn_cute_for_tests"
def _load_local_flash_attn_cute_module():
if _MODULE_NAME in sys.modules:
return sys.modules[_MODULE_NAME]
module_path = (
Path(__file__).resolve().parents[2]
/ "attention"
/ "utils"
/ "flash_attn_cute.py"
)
spec = importlib.util.spec_from_file_location(_MODULE_NAME, module_path)
if spec is None or spec.loader is None:
raise RuntimeError(f"Failed to load module spec from {module_path}")
module = importlib.util.module_from_spec(spec)
sys.modules[_MODULE_NAME] = module
spec.loader.exec_module(module)
return module
def _assert_close(
actual: torch.Tensor,
expected: torch.Tensor,
*,
dtype: torch.dtype,
is_grad: bool = False,
) -> None:
if dtype == torch.float16:
atol = 6e-3 if is_grad else 4e-3
rtol = 6e-3 if is_grad else 4e-3
else:
atol = 2e-2 if is_grad else 1e-2
rtol = 2e-2 if is_grad else 1e-2
torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
def _extract_out(x):
if isinstance(x, tuple):
return x[0]
return x
@pytest.fixture(scope="module")
def flash_attn_impls():
if not torch.cuda.is_available():
pytest.skip("CUDA is required for flash_attn.cute parity tests")
major, _ = torch.cuda.get_device_capability()
if major not in (9, 10, 11):
pytest.skip(
f"flash_attn.cute supports SM 9/10/11, got compute capability {major}"
)
try:
from flash_attn.cute import (
flash_attn_func as ref_flash_attn_func,
flash_attn_varlen_func as ref_flash_attn_varlen_func,
)
except ImportError as exc:
pytest.skip(f"flash_attn.cute is not available: {exc}")
mod = _load_local_flash_attn_cute_module()
return (
mod.flash_attn_func,
mod.flash_attn_varlen_func,
ref_flash_attn_func,
ref_flash_attn_varlen_func,
)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("causal", [False, True])
def test_flash_attn_func_parity_forward_backward(
flash_attn_impls, dtype: torch.dtype, causal: bool
):
if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported():
pytest.skip("bfloat16 is not supported on this GPU")
custom_flash_attn_func, _, ref_flash_attn_func, _ = flash_attn_impls
torch.manual_seed(0)
device = torch.device("cuda")
batch, seqlen, nheads, headdim = 2, 64, 4, 64
q_ref = torch.randn(
batch,
seqlen,
nheads,
headdim,
device=device,
dtype=dtype,
requires_grad=True,
)
k_ref = torch.randn_like(q_ref, requires_grad=True)
v_ref = torch.randn_like(q_ref, requires_grad=True)
q_test = q_ref.detach().clone().requires_grad_(True)
k_test = k_ref.detach().clone().requires_grad_(True)
v_test = v_ref.detach().clone().requires_grad_(True)
out_ref = _extract_out(ref_flash_attn_func(
q_ref,
k_ref,
v_ref,
softmax_scale=None,
causal=causal,
deterministic=False,
))
out_test = _extract_out(custom_flash_attn_func(
q_test,
k_test,
v_test,
dropout_p=0.0,
softmax_scale=None,
causal=causal,
deterministic=False,
))
_assert_close(out_test, out_ref, dtype=dtype)
dout = torch.randn_like(out_ref)
dq_ref, dk_ref, dv_ref = torch.autograd.grad(
(out_ref * dout).sum(), (q_ref, k_ref, v_ref)
)
dq_test, dk_test, dv_test = torch.autograd.grad(
(out_test * dout).sum(), (q_test, k_test, v_test)
)
_assert_close(dq_test, dq_ref, dtype=dtype, is_grad=True)
_assert_close(dk_test, dk_ref, dtype=dtype, is_grad=True)
_assert_close(dv_test, dv_ref, dtype=dtype, is_grad=True)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("causal", [False, True])
def test_flash_attn_varlen_func_parity_forward_backward(
flash_attn_impls, dtype: torch.dtype, causal: bool
):
if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported():
pytest.skip("bfloat16 is not supported on this GPU")
(
_,
custom_flash_attn_varlen_func,
_,
ref_flash_attn_varlen_func,
) = flash_attn_impls
torch.manual_seed(1)
device = torch.device("cuda")
nheads, headdim = 4, 64
seqlens = [40, 24, 16]
cu_vals = [0]
for seqlen in seqlens:
cu_vals.append(cu_vals[-1] + seqlen)
cu = torch.tensor(cu_vals, device=device, dtype=torch.int32)
total = int(cu[-1].item())
max_seqlen = max(seqlens)
q_ref = torch.randn(
total, nheads, headdim, device=device, dtype=dtype, requires_grad=True
)
k_ref = torch.randn_like(q_ref, requires_grad=True)
v_ref = torch.randn_like(q_ref, requires_grad=True)
q_test = q_ref.detach().clone().requires_grad_(True)
k_test = k_ref.detach().clone().requires_grad_(True)
v_test = v_ref.detach().clone().requires_grad_(True)
out_ref = _extract_out(ref_flash_attn_varlen_func(
q_ref,
k_ref,
v_ref,
cu_seqlens_q=cu,
cu_seqlens_k=cu,
max_seqlen_q=max_seqlen,
max_seqlen_k=max_seqlen,
softmax_scale=None,
causal=causal,
deterministic=False,
))
out_test = _extract_out(custom_flash_attn_varlen_func(
q_test,
k_test,
v_test,
cu,
cu,
max_seqlen,
max_seqlen,
dropout_p=0.0,
softmax_scale=None,
causal=causal,
deterministic=False,
))
_assert_close(out_test, out_ref, dtype=dtype)
dout = torch.randn_like(out_ref)
dq_ref, dk_ref, dv_ref = torch.autograd.grad(
(out_ref * dout).sum(), (q_ref, k_ref, v_ref)
)
dq_test, dk_test, dv_test = torch.autograd.grad(
(out_test * dout).sum(), (q_test, k_test, v_test)
)
_assert_close(dq_test, dq_ref, dtype=dtype, is_grad=True)
_assert_close(dk_test, dk_ref, dtype=dtype, is_grad=True)
_assert_close(dv_test, dv_ref, dtype=dtype, is_grad=True)
fastvideo/tests/attention/test_selector.py
-34
View File
@@ -1,34 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
import torch
from fastvideo.attention import selector
from fastvideo.platforms import AttentionBackendEnum
def test_get_attn_backend_cache_respects_forced_backend(monkeypatch):
selector._cached_get_attn_backend.cache_clear()
class _FakePlatform:
@staticmethod
def get_attn_backend_cls(selected_backend, head_size, dtype):
if selected_backend == AttentionBackendEnum.ATTN_QAT_TRAIN:
return "builtins.int"
return "builtins.str"
monkeypatch.setattr("fastvideo.platforms.current_platform", _FakePlatform())
supported_backends = (AttentionBackendEnum.FLASH_ATTN,
AttentionBackendEnum.TORCH_SDPA)
default_impl = selector.get_attn_backend(128, torch.bfloat16, supported_backends)
with selector.global_force_attn_backend_context_manager(
AttentionBackendEnum.ATTN_QAT_TRAIN):
forced_impl = selector.get_attn_backend(128, torch.bfloat16,
supported_backends)
default_impl_after = selector.get_attn_backend(128, torch.bfloat16,
supported_backends)
assert default_impl is str
assert forced_impl is int
assert default_impl_after is str
fastvideo/tests/nightly/i2v-wandb-summary.json
-1
View File
@@ -1 +0,0 @@
{"step_time":2.245914653001819,"_wandb":{"runtime":1434},"learning_rate":1e-05,"grad_norm":0.57421875,"avg_step_time":1.1814782944297622,"train_loss":0.07932619750499725,"vsa_sparsity":0,"_timestamp":1.750578625921253e+09,"validation_videos_40_steps":{"count":1,"videos":[{"size":420969,"path":"media/videos/validation_videos_40_steps_900_581ff5eae2909d3a7b36.mp4","_type":"video-file","sha256":"581ff5eae2909d3a7b362dcb24d060c006c09e4d4deb44b82f4aa697f6789ba7"}],"captions":false,"_type":"videos"},"_runtime":1434.62395329,"_step":901}
fastvideo/tests/nightly/i2v_validation_videos_40_steps_900_581ff5eae2909d3a7b36.mp4
BIN
View File
Binary file not shown.
fastvideo/tests/nightly/reference_video_1_sample_v0.mp4
BIN
View File
Binary file not shown.
fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py
-184
View File
@@ -1,184 +0,0 @@
import os
from pathlib import Path
from huggingface_hub import snapshot_download
import subprocess
import sys
import shutil
from fastvideo.tests.utils import compute_video_ssim_torchvision
# Import the training pipeline
sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
NUM_NODES = "1"
MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
# preprocessing
DATA_DIR = "data"
LOCAL_RAW_DATA_DIR = Path(os.path.join(DATA_DIR, "crush-smol"))
NUM_GPUS_PER_NODE_PREPROCESSING = "1"
PREPROCESSING_ENTRY_FILE_PATH = "fastvideo/pipelines/preprocess/v1_preprocess.py"
LOCAL_PREPROCESSED_DATA_DIR = Path(os.path.join(DATA_DIR, "crush-smol_processed_t2v"))
# training
NUM_GPUS_PER_NODE_TRAINING = "4"
TRAINING_ENTRY_FILE_PATH = "fastvideo/training/wan_distillation_pipeline.py"
LOCAL_TRAINING_DATA_DIR = os.path.join(LOCAL_PREPROCESSED_DATA_DIR, "combined_parquet_dataset")
LOCAL_VALIDATION_DATASET_FILE = "examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json"
LOCAL_OUTPUT_DIR = Path(os.path.join(DATA_DIR, "outputs"))
def download_data():
# create the data dir if it doesn't exist
data_dir = Path(DATA_DIR)
print(f"Creating data directory at {data_dir}")
os.makedirs(data_dir, exist_ok=True)
print(f"Downloading raw dataset to {LOCAL_RAW_DATA_DIR}...")
try:
result = snapshot_download(
repo_id="wlsaidhi/crush-smol-merged",
local_dir=str(LOCAL_RAW_DATA_DIR),
repo_type="dataset",
resume_download=True,
token=os.environ.get("HF_TOKEN"), # In case authentication is needed
)
print(f"Download completed successfully. Files downloaded to: {result}")
# Verify the download
if not LOCAL_RAW_DATA_DIR.exists():
raise RuntimeError(f"Download appeared to succeed but {LOCAL_RAW_DATA_DIR} does not exist")
# List downloaded files
print("Downloaded files:")
for file in LOCAL_RAW_DATA_DIR.rglob("*"):
if file.is_file():
print(f" - {file.relative_to(LOCAL_RAW_DATA_DIR)}")
except Exception as e:
print(f"Error during download: {str(e)}")
raise
def run_preprocessing():
# remove the local_preprocessed_data_dir if it exists
if LOCAL_PREPROCESSED_DATA_DIR.exists():
print(f"Removing local_preprocessed_data_dir: {LOCAL_PREPROCESSED_DATA_DIR}")
shutil.rmtree(LOCAL_PREPROCESSED_DATA_DIR)
# Run torchrun command
cmd = [
"torchrun",
"--nnodes", NUM_NODES,
"--nproc_per_node", NUM_GPUS_PER_NODE_PREPROCESSING,
PREPROCESSING_ENTRY_FILE_PATH,
"--model_path", MODEL_PATH,
"--seed", "42",
"--data_merge_path", os.path.join(LOCAL_RAW_DATA_DIR, "merge.txt"),
"--preprocess_video_batch_size", "1",
"--max_height", "480",
"--max_width", "832",
"--num_frames", "81",
"--dataloader_num_workers", "0",
"--output_dir", LOCAL_PREPROCESSED_DATA_DIR,
"--train_fps", "16",
"--samples_per_file", "1",
"--flush_frequency", "1",
"--video_length_tolerance_range", "5",
"--preprocess_task", "t2v",
]
process = subprocess.run(cmd, check=True)
def run_training():
cmd = [
"torchrun",
"--nnodes", NUM_NODES,
"--nproc_per_node", NUM_GPUS_PER_NODE_TRAINING,
TRAINING_ENTRY_FILE_PATH,
"--model_path", MODEL_PATH,
"--inference_mode", "False",
"--pretrained_model_name_or_path", MODEL_PATH,
"--data_path", LOCAL_TRAINING_DATA_DIR,
"--validation_dataset_file", LOCAL_VALIDATION_DATASET_FILE,
"--train_batch_size", "1",
"--num_latent_t", "8",
"--num_gpus", NUM_GPUS_PER_NODE_TRAINING,
"--sp_size", "1",
"--tp_size", "1",
"--hsdp_replicate_dim", "1",
"--hsdp_shard_dim", NUM_GPUS_PER_NODE_TRAINING,
"--train_sp_batch_size", "1",
"--dataloader_num_workers", "10",
"--gradient_accumulation_steps", "1",
"--max_train_steps", "501",
"--learning_rate", "2e-6",
"--fake_score_learning_rate", "2e-6",
"--mixed_precision", "bf16",
"--training_state_checkpointing_steps", "1000",
"--weight_only_checkpointing_steps", "1000",
"--validation_steps", "50",
"--validation_sampling_steps", "3",
"--log_validation",
"--checkpoints_total_limit", "3",
"--ema_start_step", "0",
"--training_cfg_rate", "0.0",
"--output_dir", LOCAL_OUTPUT_DIR,
"--tracker_project_name", "ci_wan_t2v_dmd_overfit",
"--num_height", "480",
"--num_width", "832",
"--num_frames", "81",
"--flow_shift", "8",
"--validation_guidance_scale", "6.0",
"--weight_decay", "0.01",
"--generator_update_interval", "5",
"--dmd_denoising_steps", "1000,757,522",
"--min_timestep_ratio", "0.02",
"--max_timestep_ratio", "0.98",
"--seed", "1000",
"--real_score_guidance_scale", "3.5",
"--dit_precision", "fp32",
"--max_grad_norm", "1.0",
"--enable_gradient_checkpointing_type", "full",
]
print(f"Running training with command: {cmd}")
process = subprocess.run(cmd, check=True)
def test_e2e_overfit_single_sample():
os.environ["WANDB_MODE"] = "online"
download_data()
run_preprocessing()
run_training()
reference_video_file = os.path.join(os.path.dirname(__file__), "reference_video_1_sample_v0.mp4")
print(f"reference_video_file: {reference_video_file}")
final_validation_video_file = os.path.join(LOCAL_OUTPUT_DIR, "validation_step_900_inference_steps_50_video_0.mp4")
print(f"final_validation_video_file: {final_validation_video_file}")
# Ensure both files exist
assert os.path.exists(reference_video_file), f"Reference video not found at {reference_video_file}"
assert os.path.exists(final_validation_video_file), f"Validation video not found at {final_validation_video_file}"
# Compute SSIM
mean_ssim, min_ssim, max_ssim = compute_video_ssim_torchvision(
reference_video_file,
final_validation_video_file,
use_ms_ssim=True # Using MS-SSIM for better quality assessment
)
print("\n===== SSIM Results for Step 900 Validation =====")
print(f"Mean MS-SSIM: {mean_ssim:.4f}")
print(f"Min MS-SSIM: {min_ssim:.4f}")
print(f"Max MS-SSIM: {max_ssim:.4f}")
assert max_ssim > 0.5, f"Max SSIM is below 0.5: {max_ssim}"
if __name__ == "__main__":
test_e2e_overfit_single_sample()
fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py
-173
View File
@@ -1,173 +0,0 @@
import os
from pathlib import Path
from huggingface_hub import snapshot_download
import shutil
import subprocess
import sys
from fastvideo.tests.utils import compute_video_ssim_torchvision
# Import the training pipeline
sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
NUM_NODES = "1"
MODEL_PATH = "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers"
# preprocessing
DATA_DIR = "data"
LOCAL_RAW_DATA_DIR = Path(os.path.join(DATA_DIR, "cats"))
NUM_GPUS_PER_NODE_PREPROCESSING = "1"
PREPROCESSING_ENTRY_FILE_PATH = "fastvideo/pipelines/preprocess/v1_preprocess.py"
LOCAL_PREPROCESSED_DATA_DIR = Path(os.path.join(DATA_DIR, "cats_preprocessed_data_i2v"))
# training
NUM_GPUS_PER_NODE_TRAINING = "4"
TRAINING_ENTRY_FILE_PATH = "fastvideo/training/wan_i2v_training_pipeline.py"
LOCAL_TRAINING_DATA_DIR = os.path.join(LOCAL_PREPROCESSED_DATA_DIR, "combined_parquet_dataset")
LOCAL_VALIDATION_DATASET_FILE = os.path.join(LOCAL_RAW_DATA_DIR, "validation_i2v_prompt_1_sample.json")
LOCAL_OUTPUT_DIR = Path(os.path.join(DATA_DIR, "outputs"))
def download_data():
# create the data dir if it doesn't exist
data_dir = Path(DATA_DIR)
print(f"Creating data directory at {data_dir}")
os.makedirs(data_dir, exist_ok=True)
print(f"Downloading raw dataset to {LOCAL_RAW_DATA_DIR}...")
try:
result = snapshot_download(
repo_id="wlsaidhi/cats-overfit-merged",
local_dir=str(LOCAL_RAW_DATA_DIR),
repo_type="dataset",
resume_download=True,
token=os.environ.get("HF_TOKEN"), # In case authentication is needed
)
print(f"Download completed successfully. Files downloaded to: {result}")
# Verify the download
if not LOCAL_RAW_DATA_DIR.exists():
raise RuntimeError(f"Download appeared to succeed but {LOCAL_RAW_DATA_DIR} does not exist")
# List downloaded files
print("Downloaded files:")
for file in LOCAL_RAW_DATA_DIR.rglob("*"):
if file.is_file():
print(f" - {file.relative_to(LOCAL_RAW_DATA_DIR)}")
except Exception as e:
print(f"Error during download: {str(e)}")
raise
def run_preprocessing():
# Run torchrun command
cmd = [
"torchrun",
"--nnodes", NUM_NODES,
"--nproc_per_node", NUM_GPUS_PER_NODE_PREPROCESSING,
PREPROCESSING_ENTRY_FILE_PATH,
"--model_path", MODEL_PATH,
"--data_merge_path", os.path.join(LOCAL_RAW_DATA_DIR, "merge_1_sample.txt"),
"--preprocess_video_batch_size", "1",
"--max_height", "480",
"--max_width", "832",
"--num_frames", "77",
"--dataloader_num_workers", "0",
"--output_dir", LOCAL_PREPROCESSED_DATA_DIR,
"--train_fps", "16",
"--samples_per_file", "1",
"--flush_frequency", "1",
"--video_length_tolerance_range", "5",
"--preprocess_task", "i2v",
]
process = subprocess.run(cmd, check=True)
def run_training():
cmd = [
"torchrun",
"--nnodes", NUM_NODES,
"--nproc_per_node", NUM_GPUS_PER_NODE_TRAINING,
TRAINING_ENTRY_FILE_PATH,
"--model_path", MODEL_PATH,
"--inference_mode", "False",
"--pretrained_model_name_or_path", MODEL_PATH,
"--data_path", LOCAL_TRAINING_DATA_DIR,
"--validation_dataset_file", LOCAL_VALIDATION_DATASET_FILE,
"--train_batch_size", "1",
"--num_latent_t", "8",
"--num_gpus", NUM_GPUS_PER_NODE_TRAINING,
"--sp_size", NUM_GPUS_PER_NODE_TRAINING,
"--tp_size", "1",
"--hsdp_replicate_dim", "1",
"--hsdp_shard_dim", NUM_GPUS_PER_NODE_TRAINING,
"--num_gpus", NUM_GPUS_PER_NODE_TRAINING,
"--train_sp_batch_size", "1",
"--dataloader_num_workers", "10",
"--gradient_accumulation_steps", "1",
"--max_train_steps", "901",
"--learning_rate", "1e-5",
"--mixed_precision", "bf16",
"--weight_only_checkpointing_steps", "6000",
"--training_state_checkpointing_steps", "6000",
"--validation_steps", "100",
"--validation_sampling_steps", "50",
"--log_validation",
"--checkpoints_total_limit", "3",
"--ema_start_step", "0",
"--training_cfg_rate", "0.1",
"--output_dir", LOCAL_OUTPUT_DIR,
"--tracker_project_name", "wan_i2v_finetune_overfit_ci",
"--num_height", "480",
"--num_width", "832",
"--num_frames", "81",
"--validation_guidance_scale", "6.0",
"--num_euler_timesteps", "50",
"--multi_phased_distill_schedule", "4000-1",
"--weight_decay", "0.01",
"--not_apply_cfg_solver",
"--dit_precision", "fp32",
"--max_grad_norm", "1.0",
]
print(f"Running training with command: {cmd}")
process = subprocess.run(cmd, check=True)
def test_e2e_overfit_single_sample():
os.environ["WANDB_MODE"] = "online"
download_data()
run_preprocessing()
run_training()
reference_video_file = os.path.join(os.path.dirname(__file__), "reference_video_1_sample_v0.mp4")
print(f"reference_video_file: {reference_video_file}")
final_validation_video_file = os.path.join(LOCAL_OUTPUT_DIR, "validation_step_900_inference_steps_50_video_0.mp4")
print(f"final_validation_video_file: {final_validation_video_file}")
# Ensure both files exist
assert os.path.exists(reference_video_file), f"Reference video not found at {reference_video_file}"
assert os.path.exists(final_validation_video_file), f"Validation video not found at {final_validation_video_file}"
# Compute SSIM
mean_ssim, min_ssim, max_ssim = compute_video_ssim_torchvision(
reference_video_file,
final_validation_video_file,
use_ms_ssim=True # Using MS-SSIM for better quality assessment
)
print("\n===== SSIM Results for Step 900 Validation =====")
print(f"Mean MS-SSIM: {mean_ssim:.4f}")
print(f"Min MS-SSIM: {min_ssim:.4f}")
print(f"Max MS-SSIM: {max_ssim:.4f}")
assert max_ssim > 0.5, f"Max SSIM is below 0.5: {max_ssim}"
if __name__ == "__main__":
test_e2e_overfit_single_sample()
fastvideo/tests/nightly/test_e2e_overfit_single_sample.py
-178
View File
@@ -1,178 +0,0 @@
import os
from pathlib import Path
from huggingface_hub import snapshot_download
import shutil
import subprocess
import sys
from fastvideo.tests.utils import compute_video_ssim_torchvision
# Import the training pipeline
sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
NUM_NODES = "1"
MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
# preprocessing
DATA_DIR = "data"
LOCAL_RAW_DATA_DIR = Path(os.path.join(DATA_DIR, "cats"))
NUM_GPUS_PER_NODE_PREPROCESSING = "1"
PREPROCESSING_ENTRY_FILE_PATH = "fastvideo/pipelines/preprocess/v1_preprocess.py"
LOCAL_PREPROCESSED_DATA_DIR = Path(os.path.join(DATA_DIR, "cats_preprocessed_data"))
# training
NUM_GPUS_PER_NODE_TRAINING = "4"
TRAINING_ENTRY_FILE_PATH = "fastvideo/training/wan_training_pipeline.py"
LOCAL_TRAINING_DATA_DIR = os.path.join(LOCAL_PREPROCESSED_DATA_DIR, "combined_parquet_dataset")
LOCAL_VALIDATION_DATASET_FILE = os.path.join(LOCAL_RAW_DATA_DIR, "validation_prompt_1_sample.json")
LOCAL_OUTPUT_DIR = Path(os.path.join(DATA_DIR, "outputs"))
def download_data():
# create the data dir if it doesn't exist
data_dir = Path(DATA_DIR)
print(f"Creating data directory at {data_dir}")
os.makedirs(data_dir, exist_ok=True)
print(f"Downloading raw dataset to {LOCAL_RAW_DATA_DIR}...")
try:
result = snapshot_download(
repo_id="wlsaidhi/cats-overfit-merged",
local_dir=str(LOCAL_RAW_DATA_DIR),
repo_type="dataset",
resume_download=True,
token=os.environ.get("HF_TOKEN"), # In case authentication is needed
)
print(f"Download completed successfully. Files downloaded to: {result}")
# Verify the download
if not LOCAL_RAW_DATA_DIR.exists():
raise RuntimeError(f"Download appeared to succeed but {LOCAL_RAW_DATA_DIR} does not exist")
# List downloaded files
print("Downloaded files:")
for file in LOCAL_RAW_DATA_DIR.rglob("*"):
if file.is_file():
print(f" - {file.relative_to(LOCAL_RAW_DATA_DIR)}")
except Exception as e:
print(f"Error during download: {str(e)}")
raise
def run_preprocessing():
# remove the local_preprocessed_data_dir if it exists
if LOCAL_PREPROCESSED_DATA_DIR.exists():
print(f"Removing local_preprocessed_data_dir: {LOCAL_PREPROCESSED_DATA_DIR}")
shutil.rmtree(LOCAL_PREPROCESSED_DATA_DIR)
# Run torchrun command
cmd = [
"torchrun",
"--nnodes", NUM_NODES,
"--nproc_per_node", NUM_GPUS_PER_NODE_PREPROCESSING,
PREPROCESSING_ENTRY_FILE_PATH,
"--model_path", MODEL_PATH,
"--data_merge_path", os.path.join(LOCAL_RAW_DATA_DIR, "merge_1_sample.txt"),
"--preprocess_video_batch_size", "1",
"--max_height", "480",
"--max_width", "832",
"--num_frames", "77",
"--dataloader_num_workers", "0",
"--output_dir", LOCAL_PREPROCESSED_DATA_DIR,
"--train_fps", "16",
"--samples_per_file", "1",
"--flush_frequency", "1",
"--video_length_tolerance_range", "5",
"--preprocess_task", "t2v",
]
process = subprocess.run(cmd, check=True)
def run_training():
cmd = [
"torchrun",
"--nnodes", NUM_NODES,
"--nproc_per_node", NUM_GPUS_PER_NODE_TRAINING,
TRAINING_ENTRY_FILE_PATH,
"--model_path", MODEL_PATH,
"--inference_mode", "False",
"--pretrained_model_name_or_path", MODEL_PATH,
"--data_path", LOCAL_TRAINING_DATA_DIR,
"--validation_dataset_file", LOCAL_VALIDATION_DATASET_FILE,
"--train_batch_size", "1",
"--num_latent_t", "8",
"--num_gpus", NUM_GPUS_PER_NODE_TRAINING,
"--sp_size", NUM_GPUS_PER_NODE_TRAINING,
"--tp_size", "1",
"--hsdp_replicate_dim", "1",
"--hsdp_shard_dim", NUM_GPUS_PER_NODE_TRAINING,
"--num_gpus", NUM_GPUS_PER_NODE_TRAINING,
"--train_sp_batch_size", "1",
"--dataloader_num_workers", "10",
"--gradient_accumulation_steps", "1",
"--max_train_steps", "901",
"--learning_rate", "5e-6",
"--mixed_precision", "bf16",
"--weight_only_checkpointing_steps", "6000",
"--training_state_checkpointing_steps", "6000",
"--validation_steps", "100",
"--validation_sampling_steps", "50",
"--log_validation",
"--checkpoints_total_limit", "3",
"--ema_start_step", "0",
"--training_cfg_rate", "0.0",
"--output_dir", LOCAL_OUTPUT_DIR,
"--tracker_project_name", "wan_finetune_overfit_ci",
"--num_height", "480",
"--num_width", "832",
"--num_frames", "81",
"--validation_guidance_scale", "3.0",
"--num_euler_timesteps", "50",
"--multi_phased_distill_schedule", "4000-1",
"--weight_decay", "0.01",
"--not_apply_cfg_solver",
"--dit_precision", "fp32",
"--max_grad_norm", "1.0",
]
print(f"Running training with command: {cmd}")
process = subprocess.run(cmd, check=True)
def test_e2e_overfit_single_sample():
os.environ["WANDB_MODE"] = "online"
download_data()
run_preprocessing()
run_training()
reference_video_file = os.path.join(os.path.dirname(__file__), "reference_video_1_sample_v0.mp4")
print(f"reference_video_file: {reference_video_file}")
final_validation_video_file = os.path.join(LOCAL_OUTPUT_DIR, "validation_step_900_inference_steps_50_video_0.mp4")
print(f"final_validation_video_file: {final_validation_video_file}")
# Ensure both files exist
assert os.path.exists(reference_video_file), f"Reference video not found at {reference_video_file}"
assert os.path.exists(final_validation_video_file), f"Validation video not found at {final_validation_video_file}"
# Compute SSIM
mean_ssim, min_ssim, max_ssim = compute_video_ssim_torchvision(
reference_video_file,
final_validation_video_file,
use_ms_ssim=True # Using MS-SSIM for better quality assessment
)
print("\n===== SSIM Results for Step 900 Validation =====")
print(f"Mean MS-SSIM: {mean_ssim:.4f}")
print(f"Min MS-SSIM: {min_ssim:.4f}")
print(f"Max MS-SSIM: {max_ssim:.4f}")
assert max_ssim > 0.5, f"Max SSIM is below 0.5: {max_ssim}"
if __name__ == "__main__":
test_e2e_overfit_single_sample()
fastvideo/tests/nightly/test_fp4_linear.py
-153
View File
@@ -1,153 +0,0 @@
#!/usr/bin/env python3
import math
import sys
import types
import torch
import torch.nn as nn
import torch.nn.functional as F
from fastvideo.layers.fp4linear import fp4_linear_forward
from fastvideo.layers.linear import ReplicatedLinear
def fail(msg: str):
print(f"[FAIL] {msg}")
sys.exit(1)
def metric_report(ref, test, name, cos_min, rel_l2_max):
diff = (test - ref).float()
ref32 = ref.float()
test32 = test.float()
mean_abs_diff = diff.abs().mean().item()
max_abs_diff = diff.abs().max().item()
rel_l2 = (diff.norm() / (ref32.norm() + 1e-12)).item()
cos = F.cosine_similarity(ref32.reshape(-1), test32.reshape(-1), dim=0).item()
print(f" [{name}] cos: {cos:.6f}, rel_l2: {rel_l2:.6e}, "
f"mean_abs: {mean_abs_diff:.6e}, max_abs: {max_abs_diff:.6e}")
if cos_min is not None and cos < cos_min:
fail(f"{name}: cosine similarity too low: {cos:.6f} < {cos_min}")
if rel_l2_max is not None and rel_l2 > rel_l2_max:
fail(f"{name}: relative L2 too high: {rel_l2:.6e} > {rel_l2_max}")
def run_case(case, fwd_cos_min=0.98, fwd_rel_l2_max=0.15,
bwd_cos_min=0.999, bwd_rel_l2_max=5e-3):
device = torch.device("cuda")
torch.manual_seed(0)
B = case["batch_size"]
S = case["seq_len"]
IN = case["in_features"]
OUT = case["out_features"]
NAME = case["name"]
M = B * S
print(f"\n== Case: {NAME} | M={M}, K={IN}, N={OUT} ==")
# Reference BF16 Linear
ref = nn.Linear(IN, OUT, bias=True).to(device=device, dtype=torch.bfloat16)
# FP4-forward path is installed by monkey-patching ReplicatedLinear.forward.
test_layer = ReplicatedLinear(
IN, OUT, bias=True, params_dtype=torch.float32
).to(device)
test_layer.forward = types.MethodType(fp4_linear_forward, test_layer)
# Initialize identically from the same fp32 master params
with torch.no_grad():
w32 = torch.empty(OUT, IN, device=device, dtype=torch.float32)
nn.init.kaiming_uniform_(w32, a=math.sqrt(5.0))
b32 = torch.empty(OUT, device=device, dtype=torch.float32)
nn.init.uniform_(b32, -1 / math.sqrt(IN), 1 / math.sqrt(IN))
test_layer.weight.copy_(w32)
if test_layer.bias is not None:
test_layer.bias.copy_(b32)
ref.weight.copy_(w32.to(dtype=torch.bfloat16))
if ref.bias is not None:
ref.bias.copy_(b32.to(dtype=torch.bfloat16))
# Inputs
x = torch.randn(M, IN, device=device, dtype=torch.bfloat16, requires_grad=True)
x_ref = x.clone().detach().requires_grad_(True)
# -------- Forward comparison --------
with torch.inference_mode():
y_ref = ref(x_ref)
y_fp4, _ = test_layer(x)
print("Forward metrics (FP4-forward vs BF16):")
metric_report(y_ref, y_fp4, "forward", fwd_cos_min, fwd_rel_l2_max)
# -------- Backward comparison --------
# Inject the SAME external grad_out so gradients are comparable.
grad_out = torch.randn_like(y_ref)
# Clear grads
for p in ref.parameters():
if p.grad is not None:
p.grad.zero_()
for p in test_layer.parameters():
if p.grad is not None:
p.grad.zero_()
if x.grad is not None:
x.grad.zero_()
if x_ref.grad is not None:
x_ref.grad.zero_()
# Recompute forward with grad graph (no inference_mode)
y_ref = ref(x_ref)
y_fp4, _ = test_layer(x)
# Backward with same grad_out
y_ref.backward(grad_out)
y_fp4.backward(grad_out)
# Collect grads (cast to fp32 for fair comparison)
gx_ref = x_ref.grad.float()
gw_ref = ref.weight.grad.float()
gb_ref = ref.bias.grad.float() if ref.bias is not None else None
gx_fp4 = x.grad.float()
# test_layer stores weight in fp32 → grad is fp32 already
gw_fp4 = test_layer.weight.grad.float()
gb_fp4 = test_layer.bias.grad.float() if test_layer.bias is not None else None
print("Backward metrics (grads parity with injected grad_out):")
metric_report(gx_ref, gx_fp4, "grad_x", bwd_cos_min, bwd_rel_l2_max)
metric_report(gw_ref, gw_fp4, "grad_w", bwd_cos_min, bwd_rel_l2_max)
if gb_ref is not None and gb_fp4 is not None:
metric_report(gb_ref, gb_fp4, "grad_b", bwd_cos_min, bwd_rel_l2_max)
def main():
if not torch.cuda.is_available():
print("CUDA is required.")
sys.exit(1)
# Your exact config list
configs = [
{'batch_size': 1, 'seq_len': 1, 'in_features': 256, 'out_features': 5120, 'name': 'fc_in_256x5120'},
{'batch_size': 1, 'seq_len': 1, 'in_features': 5120, 'out_features': 5120, 'name': 'fc_out_5120x5120'},
{'batch_size': 1, 'seq_len': 1, 'in_features': 5120, 'out_features': 30720,'name': 'fc_5120x30720'},
{'batch_size': 1, 'seq_len': 512,'in_features': 4096, 'out_features': 5120, 'name': 'fc_in_512x4096x5120'},
{'batch_size': 1, 'seq_len': 512,'in_features': 5120, 'out_features': 5120, 'name': 'fc_out_512x5120x5120'},
{'batch_size': 1, 'seq_len': 32768,'in_features': 5120,'out_features': 5120, 'name': 'to_q_32768x5120x5120'},
{'batch_size': 1, 'seq_len': 32768,'in_features': 5120,'out_features': 13824,'name': 'ffn_fc_in_32768x5120x13824'},
{'batch_size': 1, 'seq_len': 32768,'in_features': 13824,'out_features': 5120, 'name': 'ffn_fc_out_32768x13824x5120'},
]
# Tolerances: forward is approximate (FP4), backward should match closely
fwd_cos_min = 0.98
fwd_rel_l2_max = 0.15
bwd_cos_min = 0.999
bwd_rel_l2_max = 5e-3
for case in configs:
run_case(case, fwd_cos_min, fwd_rel_l2_max, bwd_cos_min, bwd_rel_l2_max)
torch.cuda.synchronize()
print("\n[OK] All forward & backward checks passed for all cases.")
if __name__ == "__main__":
main()
fastvideo/tests/nightly/test_sage3_with_16bit_bwd.py
-137
View File
@@ -1,137 +0,0 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
import sys
import math
import torch
import torch.nn.functional as F
# ---- Import function-under-test ----
# Make sure this import path matches your project layout.
# If it's in the same file, remove this import and keep the definition above.
from fastvideo.attention.backends.sage_attn3 import attn_forward_4bit_fwd_16bit_bwd
from fastvideo.attention.backends.sage_attn3 import sageattn_blackwell_with_16bit_bwd
# ---- Import FlashAttention forward (public) ----
# We DO NOT use _wrapped_* internals. Prefer the top-level alias; fall back if needed.
try:
from flash_attn import flash_attn_func as flash_attn
except Exception:
try:
from flash_attn.flash_attn_interface import flash_attn_func as flash_attn
except Exception as e:
print("Could not import flash_attn function:", e)
sys.exit(1)
def fail(msg: str):
print(f"[FAIL] {msg}")
sys.exit(1)
def metric_report(ref, test, name, cos_min=None, rel_l2_max=None):
ref32 = ref.float()
test32 = test.float()
diff = test32 - ref32
mean_abs = diff.abs().mean().item()
max_abs = diff.abs().max().item()
rel_l2 = (diff.norm() / (ref32.norm() + 1e-12)).item()
cos = F.cosine_similarity(ref32.reshape(-1), test32.reshape(-1), dim=0).item()
print(f" [{name}] cos={cos:.6f} rel_l2={rel_l2:.6e} mean_abs={mean_abs:.6e} max_abs={max_abs:.6e}")
if cos_min is not None and cos < cos_min:
fail(f"{name}: cosine too low: {cos:.6f} < {cos_min}")
if rel_l2_max is not None and rel_l2 > rel_l2_max:
fail(f"{name}: relative L2 too high: {rel_l2:.6e} > {rel_l2_max}")
def run_case(case, fwd_cos_min=0.95, fwd_rel_l2_max=0.25, bwd_cos_min=0.95, bwd_rel_l2_max=0.25):
if not torch.cuda.is_available():
print("CUDA is required.")
sys.exit(1)
device = torch.device("cuda")
torch.manual_seed(case.get("seed", 0))
B = case["B"]
L = case["L"]
H = case["H"]
D = case["D"]
dtype = case["dtype"]
causal = case["causal"]
name = case["name"]
print(f"\n== Case: {name} | B={B}, L={L}, H={H}, D={D}, dtype={str(dtype).split('.')[-1]}, causal={causal} ==")
# Inputs in BLHD (flash_attn expects BLHD)
q = torch.randn(B, L, H, D, device=device, dtype=dtype, requires_grad=True)
k = torch.randn(B, L, H, D, device=device, dtype=dtype, requires_grad=True)
v = torch.randn(B, L, H, D, device=device, dtype=dtype, requires_grad=True)
softmax_scale = 1.0 / math.sqrt(D)
# ----- Forward -----
# Our kernel (forward uses sage, backward will use FA inside the autograd fn)
y_sage = sageattn_blackwell_with_16bit_bwd(q, k, v, is_causal=causal)
# Sanity checks
if y_sage.shape != (B, L, H, D):
fail(f"Output shape mismatch: got {tuple(y_sage.shape)}, expected {(B, L, H, D)}")
if not torch.isfinite(y_sage).all():
fail("Non-finite values detected in forward output")
# FlashAttention reference forward (public function)
with torch.inference_mode():
y_fa = attn_forward_4bit_fwd_16bit_bwd(q, k, v, is_causal=causal)
print("Forward numeric check vs FlashAttention:")
metric_report(y_fa, y_sage, "forward", cos_min=fwd_cos_min, rel_l2_max=fwd_rel_l2_max)
# ----- Backward -----
# Inject the SAME grad_out for both models so gradients are comparable
grad_out = torch.randn_like(y_sage)
# Keep separate copies with grad for the FA path
q_ref = q.detach().clone().requires_grad_(True)
k_ref = k.detach().clone().requires_grad_(True)
v_ref = v.detach().clone().requires_grad_(True)
# Backprop through our function
# (It recomputes FA fwd+bwd internally in 16-bit per your implementation)
for t in (q, k, v):
if t.grad is not None:
t.grad.zero_()
y_sage.backward(grad_out)
dq_sage, dk_sage, dv_sage = q.grad.detach(), k.grad.detach(), v.grad.detach()
# Backprop through FA public forward
for t in (q_ref, k_ref, v_ref):
if t.grad is not None:
t.grad.zero_()
y_ref = attn_forward_4bit_fwd_16bit_bwd(q_ref, k_ref, v_ref, is_causal=causal)
y_ref.backward(grad_out)
dq_ref, dk_ref, dv_ref = q_ref.grad.detach(), k_ref.grad.detach(), v_ref.grad.detach()
print("Backward gradient checks vs FlashAttention:")
metric_report(dq_ref, dq_sage, "grad_q", cos_min=bwd_cos_min, rel_l2_max=bwd_rel_l2_max)
metric_report(dk_ref, dk_sage, "grad_k", cos_min=bwd_cos_min, rel_l2_max=bwd_rel_l2_max)
metric_report(dv_ref, dv_sage, "grad_v", cos_min=bwd_cos_min, rel_l2_max=bwd_rel_l2_max)
def main():
# Keep sizes moderate so it runs quickly; adjust as you like.
cases = [
{"name": "bf16_noncausal_L16384", "B": 1, "L": 1024, "H": 8, "D": 128, "dtype": torch.bfloat16, "causal": False},
{"name": "bf16_noncausal_L2048", "B": 1, "L": 9360, "H": 8, "D": 128, "dtype": torch.bfloat16, "causal": False},
]
for case in cases:
run_case(case)
torch.cuda.synchronize()
print("\n[OK] All tests passed.")
if __name__ == "__main__":
main()
fastvideo/tests/nightly/test_torch_fp4_qk.py
-85
View File
@@ -1,85 +0,0 @@
#!/usr/bin/env python3
import math
import sys
import torch
import torch.nn.functional as F
from fastvideo.attention.backends.sage_attn3 import fp_16_qk, sage_attn_qk_torch
def fail(msg: str):
print(f"[FAIL] {msg}")
sys.exit(1)
def metric_report(ref, test, name, cos_min=None, rel_l2_max=None):
diff = (test - ref).float()
ref32 = ref.float()
test32 = test.float()
mean_abs_diff = diff.abs().mean().item()
max_abs_diff = diff.abs().max().item()
rel_l2 = (diff.norm() / (ref32.norm() + 1e-12)).item()
cos = F.cosine_similarity(ref32.reshape(-1), test32.reshape(-1), dim=0).item()
print(f" [{name}] cos: {cos:.6f}, rel_l2: {rel_l2:.6e}, "
f"mean_abs: {mean_abs_diff:.6e}, max_abs: {max_abs_diff:.6e}")
if cos_min is not None and cos < cos_min:
fail(f"{name}: cosine similarity too low: {cos:.6f} < {cos_min}")
if rel_l2_max is not None and rel_l2 > rel_l2_max:
fail(f"{name}: relative L2 too high: {rel_l2:.6e} > {rel_l2_max}")
def run_case(case, fwd_cos_min=0.98, fwd_rel_l2_max=2.0e-1):
device = torch.device("cuda")
torch.manual_seed(0)
B, H, L, D = case["B"], case["H"], case["L"], case["D"]
NAME = case["name"]
assert B == 1, "sage_attn_qk_torch requires B == 1"
print(f"\n== Case: {NAME} | B={B}, H={H}, L={L}, D={D} ==")
dtype = torch.bfloat16
q = torch.randn(B, H, L, D, device=device, dtype=dtype)
k = torch.randn(B, H, L, D, device=device, dtype=dtype)
with torch.inference_mode():
# Reference computes q @ k^T / sqrt(D)
p_ref = fp_16_qk(q, k.transpose(-1, -2)) # [B,H,L,L]
p_sage = sage_attn_qk_torch(q, k, False) # [B,H,L,L]
print("Forward metrics (sage_attn FP4 vs fp16/bf16 reference):")
metric_report(p_ref, p_sage, "forward", fwd_cos_min, fwd_rel_l2_max)
def main():
try:
_ = fp_16_qk
_ = sage_attn_qk_torch
except NameError:
print("Please import/define fp_16_qk and sage_attn_qk_torch before running.")
sys.exit(1)
if not torch.cuda.is_available():
print("CUDA is required.")
sys.exit(1)
configs = [
{"B": 1, "H": 4, "L": 256, "D": 128, "name": "H4_L256_D128"},
{"B": 1, "H": 8, "L": 512, "D": 128, "name": "H8_L512_D128"},
{"B": 1, "H": 8, "L": 256, "D": 256, "name": "H8_L256_D256"},
{"B": 1, "H": 16, "L": 384, "D": 128, "name": "H16_L384_D128"},
]
fwd_cos_min = 0.98
fwd_rel_l2_max = 2.0e-1 # relax for FP4 probabilities
for case in configs:
run_case(case, fwd_cos_min=fwd_cos_min, fwd_rel_l2_max=fwd_rel_l2_max)
torch.cuda.synchronize()
print("\n[OK] All forward diff checks passed for all cases.")
if __name__ == "__main__":
main()
fastvideo/tests/ops/quantization/test_absmax_fp8.py
-94
View File
@@ -1,94 +0,0 @@
import unittest
import torch
import torch.nn as nn
from torch.testing import assert_close
from fastvideo.layers.quantization.absmax_fp8 import (
AbsMaxFP8LinearMethod,
AbsMaxFP8MergedParameter,
AbsMaxFP8Parameter,
)
from fastvideo.models.utils import set_weight_attrs
class TestAbsMaxFP8LinearMethod(unittest.TestCase):
def test_convert_scale_none(self):
method = AbsMaxFP8LinearMethod()
scale = method._convert_scale(None)
self.assertIsInstance(scale, AbsMaxFP8Parameter)
self.assertEqual(scale.dtype, torch.float32)
assert_close(scale, torch.tensor([1.0], dtype=torch.float32))
def test_convert_scale_scalar(self):
method = AbsMaxFP8LinearMethod()
scale = method._convert_scale(2.5)
self.assertIsInstance(scale, AbsMaxFP8Parameter)
self.assertEqual(scale.dtype, torch.float32)
assert_close(scale, torch.tensor([2.5], dtype=torch.float32))
def test_convert_scale_rejects_non_float32(self):
method = AbsMaxFP8LinearMethod()
scale = torch.tensor([1.0], dtype=torch.float16)
with self.assertRaisesRegex(NotImplementedError, "float32"):
method._convert_scale(scale)
def test_create_weights_rejects_invalid_dtype(self):
method = AbsMaxFP8LinearMethod()
layer = nn.Module()
with self.assertRaisesRegex(AssertionError, "only supports"):
method.create_weights(
layer=layer,
input_size_per_partition=2,
output_partition_sizes=[3],
input_size=2,
output_size=3,
params_dtype=torch.float32,
)
def test_absmax_fp8_parameter_weight_loader(self):
param = AbsMaxFP8Parameter(torch.zeros(1), requires_grad=False)
param.weight_loader(param, torch.tensor(3.0))
assert_close(param, torch.tensor([3.0]))
def test_absmax_fp8_merged_parameter_weight_loader(self):
method = AbsMaxFP8LinearMethod()
output_partition_sizes = [2, 3, 4]
param = method._merged_placeholder(output_partition_sizes)
self.assertIsInstance(param, AbsMaxFP8MergedParameter)
param.weight_loader(param, torch.tensor(7.0), share_id="k")
expected = torch.ones(sum(output_partition_sizes), dtype=torch.float32)
expected[2:5] = 7.0
assert_close(param, expected)
def test_absmax_fp8_merged_parameter_rejects_invalid_share_id(self):
method = AbsMaxFP8LinearMethod()
param = method._merged_placeholder([2, 2, 2])
with self.assertRaisesRegex(ValueError, "requires share_id"):
param.weight_loader(param, torch.tensor(1.0), share_id="bad")
def test_apply_matches_linear(self):
method = AbsMaxFP8LinearMethod()
layer = nn.Module()
method.create_weights(
layer=layer,
input_size_per_partition=3,
output_partition_sizes=[2],
input_size=3,
output_size=2,
params_dtype=torch.float16,
)
weight_fp16 = torch.tensor(
[[1.0, -2.0, 3.0], [4.0, 0.5, -1.5]], dtype=torch.float16
)
layer.weight.data = weight_fp16.to(dtype=torch.float8_e4m3fn)
layer.scale_weight.data = torch.tensor([2.0, 3.0], dtype=torch.float32)
layer.scale_input.data = torch.tensor([4.0], dtype=torch.float32)
x = torch.tensor([[1.0, 2.0, -1.0]], dtype=torch.float16)
expected = torch.nn.functional.linear(
x * layer.scale_input.data.to(dtype=torch.float16),
weight_fp16
* layer.scale_weight.data.to(dtype=torch.float16).unsqueeze(1),
).to(dtype=torch.float16)
output = method.apply(layer, x, bias=None)
assert_close(output, expected)
fastvideo/tests/ops/test_linear.py
-39
View File
@@ -1,39 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
import torch
from fastvideo.layers.linear import ReplicatedLinear
def test_replicated_linear_shape_tracking_disabled_by_default():
ReplicatedLinear.reset_shape_tracking()
ReplicatedLinear.enable_shape_tracking = False
layer = ReplicatedLinear(4, 6, bias=False)
x = torch.randn(2, 4)
layer(x)
assert ReplicatedLinear.get_shape_mapping() == {}
def test_replicated_linear_shape_tracking_records_shapes_when_enabled():
ReplicatedLinear.reset_shape_tracking()
ReplicatedLinear.enable_shape_tracking = True
try:
layer = ReplicatedLinear(4, 6, bias=False, prefix="test.linear")
x = torch.randn(2, 4)
output, _ = layer(x)
output_again, _ = layer(x)
shape_mapping = ReplicatedLinear.get_shape_mapping()
shape_key = (x.shape, output.shape)
assert output.shape == output_again.shape
assert shape_key in shape_mapping
assert shape_mapping[shape_key] == ["ReplicatedLinear"]
finally:
ReplicatedLinear.enable_shape_tracking = False
ReplicatedLinear.reset_shape_tracking()