diff --git a/fastvideo/tests/attention/test_flash_attn_cute_custom_op.py b/fastvideo/tests/attention/test_flash_attn_cute_custom_op.py deleted file mode 100644 index a3f7bb476..000000000 --- a/fastvideo/tests/attention/test_flash_attn_cute_custom_op.py +++ /dev/null @@ -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) diff --git a/fastvideo/tests/attention/test_selector.py b/fastvideo/tests/attention/test_selector.py deleted file mode 100644 index 6685138ae..000000000 --- a/fastvideo/tests/attention/test_selector.py +++ /dev/null @@ -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 diff --git a/fastvideo/tests/nightly/i2v-wandb-summary.json b/fastvideo/tests/nightly/i2v-wandb-summary.json deleted file mode 100644 index 511c05de6..000000000 --- a/fastvideo/tests/nightly/i2v-wandb-summary.json +++ /dev/null @@ -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} \ No newline at end of file diff --git a/fastvideo/tests/nightly/i2v_validation_videos_40_steps_900_581ff5eae2909d3a7b36.mp4 b/fastvideo/tests/nightly/i2v_validation_videos_40_steps_900_581ff5eae2909d3a7b36.mp4 deleted file mode 100644 index e5ff30261..000000000 Binary files a/fastvideo/tests/nightly/i2v_validation_videos_40_steps_900_581ff5eae2909d3a7b36.mp4 and /dev/null differ diff --git a/fastvideo/tests/nightly/reference_video_1_sample_v0.mp4 b/fastvideo/tests/nightly/reference_video_1_sample_v0.mp4 deleted file mode 100644 index 6918b8bac..000000000 Binary files a/fastvideo/tests/nightly/reference_video_1_sample_v0.mp4 and /dev/null differ diff --git a/fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py b/fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py deleted file mode 100644 index 1adf817f3..000000000 --- a/fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py +++ /dev/null @@ -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() diff --git a/fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py b/fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py deleted file mode 100644 index 88c409f40..000000000 --- a/fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py +++ /dev/null @@ -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() diff --git a/fastvideo/tests/nightly/test_e2e_overfit_single_sample.py b/fastvideo/tests/nightly/test_e2e_overfit_single_sample.py deleted file mode 100644 index 93f44eeb3..000000000 --- a/fastvideo/tests/nightly/test_e2e_overfit_single_sample.py +++ /dev/null @@ -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() diff --git a/fastvideo/tests/nightly/test_fp4_linear.py b/fastvideo/tests/nightly/test_fp4_linear.py deleted file mode 100644 index a3c0a3f75..000000000 --- a/fastvideo/tests/nightly/test_fp4_linear.py +++ /dev/null @@ -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() diff --git a/fastvideo/tests/nightly/test_sage3_with_16bit_bwd.py b/fastvideo/tests/nightly/test_sage3_with_16bit_bwd.py deleted file mode 100644 index 6ab2080d3..000000000 --- a/fastvideo/tests/nightly/test_sage3_with_16bit_bwd.py +++ /dev/null @@ -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() diff --git a/fastvideo/tests/nightly/test_torch_fp4_qk.py b/fastvideo/tests/nightly/test_torch_fp4_qk.py deleted file mode 100644 index 3493c3dc0..000000000 --- a/fastvideo/tests/nightly/test_torch_fp4_qk.py +++ /dev/null @@ -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() \ No newline at end of file diff --git a/fastvideo/tests/ops/quantization/test_absmax_fp8.py b/fastvideo/tests/ops/quantization/test_absmax_fp8.py deleted file mode 100644 index bed3439c9..000000000 --- a/fastvideo/tests/ops/quantization/test_absmax_fp8.py +++ /dev/null @@ -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) \ No newline at end of file diff --git a/fastvideo/tests/ops/test_linear.py b/fastvideo/tests/ops/test_linear.py deleted file mode 100644 index fb3f1d8be..000000000 --- a/fastvideo/tests/ops/test_linear.py +++ /dev/null @@ -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()