commit 4df75a30cd8e1338cac09f724ae2cf28035dcd37 Author: wehub-resource-sync Date: Mon Jul 13 12:31:40 2026 +0800 chore: import upstream snapshot with attribution diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..3656a21 --- /dev/null +++ b/.gitignore @@ -0,0 +1,41 @@ +.DS_Store + +__pycache__/ +*.py[cod] +*$py.class +.pytest_cache/ +.mypy_cache/ +.ruff_cache/ + +longlive_models/ +wan_models/ +videos/ +logs/ +outputs/ +wandb/ +debug_save/ +pretrained/ +.codex/ +.claude/ + +# Local reports and generated release artifacts +*.html +*.log +*.out +*.err +*.mp4 +*.pt +*.pth + +# Optimization experiment records +agent/logs/ +agent/runs/ +agent/profiles/ + +# Native and build artifacts +*.so +*.egg-info/ +build/ +dist/ +fouroversix/build/ +utils/kernel/build/ diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..57534a7 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,141 @@ + +## LongLive OSS Contribution Rules + +#### Issue Tracking + +* All enhancement, bugfix, or change requests must begin with the creation of a [LongLive Issue Request](https://github.com/nvidia/LongLive/issues). + * The issue request must be reviewed by LongLive engineers and approved prior to code review. + + +#### Coding Guidelines + +- All source code contributions must strictly adhere to the [LongLive Coding Guidelines](CODING-GUIDELINES.md). + +- In addition, please follow the existing conventions in the relevant file, submodule, module, and project when you add new code or when you extend/fix existing functionality. + +- To maintain consistency in code formatting and style, you should also run `clang-format` on the modified sources with the provided configuration file. This applies LongLive code formatting rules to: + - class, function/method, and variable/field naming + - comment style + - indentation + - line length + +- Format git changes: + ```bash + # Commit ID is optional - if unspecified, run format on staged changes. + git-clang-format --style file [commit ID/reference] + ``` + +- Format individual source files: + ```bash + # -style=file : Obtain the formatting rules from .clang-format + # -i : In-place modification of the processed file + clang-format -style=file -i -fallback-style=none + ``` + +- Format entire codebase (for project maintainers only): + ```bash + find samples plugin -iname *.h -o -iname *.c -o -iname *.cpp -o -iname *.hpp \ + | xargs clang-format -style=file -i -fallback-style=none + ``` + +- Avoid introducing unnecessary complexity into existing code so that maintainability and readability are preserved. + +- Try to keep pull requests (PRs) as concise as possible: + - Avoid committing commented-out code. + - Wherever possible, each PR should address a single concern. If there are several otherwise-unrelated things that should be fixed to reach a desired endpoint, our recommendation is to open several PRs and indicate the dependencies in the description. The more complex the changes are in a single PR, the more time it will take to review those changes. + +- Write commit titles using imperative mood and [these rules](https://chris.beams.io/posts/git-commit/), and reference the Issue number corresponding to the PR. Following is the recommended format for commit texts: +``` +# - + + +``` + +- Ensure that the build log is clean, meaning no warnings or errors should be present. + +- Ensure that all `sample_*` tests pass prior to submitting your code. + +- All OSS components must contain accompanying documentation (READMEs) describing the functionality, dependencies, and known issues. + + - See `README.md` for existing samples and plugins for reference. + +- All OSS components must have an accompanying test. + + - If introducing a new component, such as a plugin, provide a test sample to verify the functionality. + +- To add or disable functionality: + - Add a CMake option with a default value that matches the existing behavior. + - Where entire files can be included/excluded based on the value of this option, selectively include/exclude the relevant files from compilation by modifying `CMakeLists.txt` rather than using `#if` guards around the entire body of each file. + - Where the functionality involves minor changes to existing files, use `#if` guards. + +- Make sure that you can contribute your work to open source (no license and/or patent conflict is introduced by your code). You will need to [`sign`](#signing-your-work) your commit. + +- Thanks in advance for your patience as we review your contributions; we do appreciate them! + + +#### Pull Requests +Developer workflow for code contributions is as follows: + +1. Developers must first [fork](https://help.github.com/en/articles/fork-a-repo) the [upstream](https://github.com/nvidia/LongLive) LongLive OSS repository. + +2. Git clone the forked repository and push changes to the personal fork. + + ```bash +git clone https://github.com/YOUR_USERNAME/YOUR_FORK.git LongLive +# Checkout the targeted branch and commit changes +# Push the commits to a branch on the fork (remote). +git push -u origin : + ``` + +3. Once the code changes are staged on the fork and ready for review, a [Pull Request](https://help.github.com/en/articles/about-pull-requests) (PR) can be [requested](https://help.github.com/en/articles/creating-a-pull-request) to merge the changes from a branch of the fork into a selected branch of upstream. + * Exercise caution when selecting the source and target branches for the PR. + Note that versioned releases of LongLive OSS are posted to `release/` branches of the upstream repo. + * Creation of a PR creation kicks off the code review process. + * Atleast one LongLive engineer will be assigned for the review. + * While under review, mark your PRs as work-in-progress by prefixing the PR title with [WIP]. + +4. Since there is no CI/CD process in place yet, the PR will be accepted and the corresponding issue closed only after adequate testing has been completed, manually, by the developer and/or LongLive engineer reviewing the code. + + +#### Signing Your Work + +* We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license. + + * Any contribution which contains commits that are not Signed-Off will not be accepted. + +* To sign off on a commit you simply use the `--signoff` (or `-s`) option when committing your changes: + ```bash + $ git commit -s -m "Add cool feature." + ``` + This will append the following to your commit message: + ``` + Signed-off-by: Your Name + ``` + +* Full text of the DCO: + + ``` + Developer Certificate of Origin + Version 1.1 + + Copyright (C) 2004, 2006 The Linux Foundation and its contributors. + 1 Letterman Drive + Suite D4700 + San Francisco, CA, 94129 + + Everyone is permitted to copy and distribute verbatim copies of this license document, but changing it is not allowed. + ``` + + ``` + Developer's Certificate of Origin 1.1 + + By making a contribution to this project, I certify that: + + (a) The contribution was created in whole or in part by me and I have the right to submit it under the open source license indicated in the file; or + + (b) The contribution is based upon previous work that, to the best of my knowledge, is covered under an appropriate open source license and I have the right under that license to submit that work with modifications, whether created in whole or in part by me, under the same open source license (unless I am permitted to submit under a different license), as indicated in the file; or + + (c) The contribution was provided directly to me by some other person who certified (a), (b) or (c) and I have not modified it. + + (d) I understand and agree that this project and the contribution are public and that a record of the contribution (including all personal information I submit with it, including my sign-off) is maintained indefinitely and may be redistributed consistent with this project or the open source license(s) involved. + ``` diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..261eeb9 --- /dev/null +++ b/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/README.md b/README.md new file mode 100644 index 0000000..4a10204 --- /dev/null +++ b/README.md @@ -0,0 +1,282 @@ +

+ LongLive2.0 logo +

+ +# 🎬 LongLive 2.0: An NVFP4 Parallel Infrastructure for Long Video Generation + +[![Paper](https://img.shields.io/badge/Paper-LongLive_2.0-brown)](https://arxiv.org/abs/2605.18739) +[![Paper](https://img.shields.io/badge/Paper-LongLive_1.0-orange)](https://github.com/NVlabs/LongLive/tree/v1.0) +[![Paper](https://img.shields.io/badge/Paper-LongLive_RAG-yellow)](https://github.com/qixinhu11/LongLive-RAG) +[![Video](https://img.shields.io/badge/YouTube-Video-red)](https://www.youtube.com/watch?v=7oQALy32fiU) +[![Code](https://img.shields.io/badge/GitHub-Code-blue)](https://github.com/NVlabs/LongLive) +[![Demo](https://img.shields.io/badge/Demo-Page-green)](https://nvlabs.github.io/LongLive/LongLive2/) +[![Docs](https://img.shields.io/badge/Full-Documentation-brightgreen)](https://nvlabs.github.io/LongLive/LongLive2/docs/) + + +
+ + + +[![Watch the video](assets/longlive2/first-video-frame.png)](https://www.youtube.com/watch?v=7oQALy32fiU) + +
+ +## 💡 TLDR: Infra with NVFP4 and parallelism for both training and inference + +

+ LongLive2.0 teaser +

+ +## News +- 🔥 [2026.07.08] We support FP8 inference on LongLive 2.0. Please refer to [here](https://github.com/NVlabs/LongLive#fp8-ptq). +- 🔥 [2026.06.01] We released [LongLive-RAG](https://github.com/qixinhu11/LongLive-RAG), a general retrieval-augmented framework for long video gen. +- 🔥 [2026.05.30] LongLive2.0 now supports I2V AR teacher-forcing training and I2V DMD distillation for Wan2.2-TI2V-5B. +- ⚡ [2026.05.25] We optimized the NVFP4 inference path with fused Triton RoPE/adaLN kernels, reduced KV-cache synchronization overhead, in-place quantized KV-cache updates, faster FP4 KV dequantization, pinned VAE transfers, and safer LoRA-before-quantization setup, improving overall throughput by **18.6%**. +- 🔥 [2026.05.13] We release **LongLive 2.0**, infra with NVFP4, parallelism and multi-shot for AR training, DMD distillation, and inference (⚡45.7 FPS). The original LongLive 1.0 is now in the [v1.0](https://github.com/NVlabs/LongLive/tree/v1.0) branch. +- 🔥 [2026.04.12] LongLive supports kv cache compression with [TriAttention](https://github.com/WeianMao/triattention), with 50% KV reduction and no quality drop. Check it [here](https://github.com/WeianMao/triattention/tree/main/longlive) +- 🎉 [2026.1.27] LongLive is accepted by **ICLR-2026**. +- 🔥 [2026.1.11] LongLive supports adapting LongLive's original RoPE into KV-cache relative RoPE and generates infinite long videos! +- 🔥 [2025.11.3] We implement LongLive on linear attention model [SANA-Video](https://nvlabs.github.io/Sana/Video/)! Now SANA-Video can generate 60s interactive videos in real-time. +- 🔥 [2025.9.29] We release [Paper](https://arxiv.org/abs/2509.22622), this GitHub repo [LongLive](https://github.com/NVlabs/LongLive) with all training and inference code, the model weight [LongLive-1.3B](https://huggingface.co/Efficient-Large-Model/LongLive-1.3B), and demo page [Website](https://nvlabs.github.io/LongLive). + +## Introduction + +**LongLive 1.0**: Real-time Interactive Long Video Generation. [You can find it here](https://github.com/NVlabs/LongLive/tree/v1.0) in our V1.0 branch. + +**LongLive 2.0**: an NVFP4 Parallel Infrastructure for Long Video Generation +- For training, it supports + - [x] Balanced sequence parallel for T2V/I2V AR training (teacher-forcing). + - [x] T2V/I2V AR training on multi-shot (or single-shot) videos. + - [x] NVFP4 (or BF16) for both AR training and few-step distillation. +- For inference, it supports + - [x] NVFP4 inference (W4A4) and NVFP4 KV Cache. + - [x] TorchAO FP8 PTQ inference (W8A8) from the BF16 checkpoint. + - [x] Multi-shot attention sink. + - [x] Sequence parallel inference. + - [x] Async decoding. + + +

+ LongLive2.0 framework overview +

+ + +**LongLive 1.0**: Real-time Interactive Long Video Generation. It accepts sequential user prompts and generates corresponding videos in real time, enabling user-guided long video generation. The key insights are attention sink, KV-recache, and streaming long tuning. + + +

+ LongLive1.0 framework overview +

+ +## Getting Started +- [Full Documentation](https://nvlabs.github.io/LongLive/LongLive2/docs/) +- [Installation](https://nvlabs.github.io/LongLive/LongLive2/docs/#installation) +- [NVFP4 Setup](https://nvlabs.github.io/LongLive/LongLive2/docs/#nvfp4-installation) +- [Training Modes](https://nvlabs.github.io/LongLive/LongLive2/docs/#training) +- [Inference](https://nvlabs.github.io/LongLive/LongLive2/docs/#inference) +- [Data Organization](https://nvlabs.github.io/LongLive/LongLive2/docs/#training-data) + + +The default git clone fetches objects from all branches, including our demopage branch, which contains large assets. For normal use, only the main branch is needed. Please clone only main with: + +```git clone --single-branch --branch main --depth 1 https://github.com/NVlabs/LongLive.git``` + +### Quick Start + +#### BF16 + +```python +import torch +from omegaconf import OmegaConf + +from pipeline import CausalDiffusionInferencePipeline +from utils.config import normalize_config +from utils.inference_utils import ( + load_generator_checkpoint, + place_vae_for_streaming, + prepare_single_prompt_inputs, + save_video, +) + +prompt = "A compact silver robot walks through a clean robotics lab." +merged_checkpoint_path = "LongLive-2.0-5B/model_bf16.pt" + +config = normalize_config(OmegaConf.load("configs/inference.yaml")) +device = torch.device("cuda") + +torch.set_grad_enabled(False) +pipe = CausalDiffusionInferencePipeline(config, device=device) +load_generator_checkpoint(pipe.generator, merged_checkpoint_path) +pipe = pipe.to(device=device, dtype=torch.bfloat16) +place_vae_for_streaming(pipe, config) # honor streaming_vae + vae_device when set +pipe.generator.model.eval().requires_grad_(False) + +noise, prompts = prepare_single_prompt_inputs(config, prompt, device) +video = pipe.inference(noise=noise, text_prompts=prompts) +save_video(video[0], "videos/quickstart/sample.mp4", fps=24) +``` + +`place_vae_for_streaming` is a no-op unless `inference.streaming_vae` is true and `inference.vae_device` is set, so toggling streaming-pipeline decode in your yaml is enough — the script does not need to change. + +#### FP8 PTQ + +Download `model_bf16.pt` from +[`Efficient-Large-Model/LongLive-2.0-5B`](https://huggingface.co/Efficient-Large-Model/LongLive-2.0-5B), +set `checkpoints.generator_ckpt` in `configs/fp8/inference_fp8.yaml`, and run: + +```bash +python inference.py --config_path configs/fp8/inference_fp8.yaml +``` + +This loads the BF16 generator, applies TorchAO row-wise dynamic FP8 W8A8 PTQ, +and then enables the existing `torch.compile` path. With the provided 5B model, +300 eligible core Linear layers use FP8; six small conditioning/output +projections stay in BF16 for stability and to avoid FP8 overhead. + +The validated stack is Python 3.10, PyTorch 2.8.0+cu128, and TorchAO 0.13.0 on +H100 (SM90); compute capability 8.9 or newer is required. The supplied config +uses `torch_compile: auto`. Its `max-autotune` warm-up can take several minutes +while guard/shape variants are compiled, so use repeated inference and exclude +all compile/warm-up samples when measuring steady-state performance. Set +`torch_compile: false` for a short eager-mode smoke test. + +The supplied config uses the single 8-latent-frame block validated on H100. +Longer generation introduces additional KV-cache shapes and may trigger more +compilation or eager fallback; validate the intended frame count before +benchmarking or deployment. + +#### NVFP4 + +Point `checkpoints.generator_ckpt` in `configs/nvfp4/inference_nvfp4.yaml` at the downloaded checkpoint and set `model_quant_use_transformer_engine` according to the backend you are using: + +- TransformerEngine checkpoint (`model_te.pt`): `model_quant_use_transformer_engine: true` +- FourOverSix checkpoint (`model_4o6.pt`): `model_quant_use_transformer_engine: false` + +`setup_nvfp4_pipeline` handles checkpoint loading, NVFP4 module wrapping, weight materialization, dtype/device placement, and the streaming-pipeline VAE relocation for both backends — the bf16 `pipe.to(...)` shortcut is unsafe here because it would cast the quantized buffers. + +```python +import torch +from omegaconf import OmegaConf + +from pipeline import CausalDiffusionInferencePipeline +from utils.config import normalize_config +from utils.inference_utils import prepare_single_prompt_inputs, save_video, setup_nvfp4_pipeline + +prompt = "A compact silver robot walks through a clean robotics lab." + +config = normalize_config(OmegaConf.load("configs/nvfp4/inference_nvfp4.yaml")) +device = torch.device("cuda") + +torch.set_grad_enabled(False) +pipe = CausalDiffusionInferencePipeline(config, device=device) +setup_nvfp4_pipeline(pipe, config, device) +pipe.generator.model.eval().requires_grad_(False) + +noise, prompts = prepare_single_prompt_inputs(config, prompt, device) +video = pipe.inference(noise=noise, text_prompts=prompts) +save_video(video[0], "videos/quickstart/sample_nvfp4.mp4", fps=24) +``` + +## Training Modes + +LongLive2.0 supports both T2V and I2V training. Each modality follows the same two-stage recipe: AR teacher-forcing training first, then DMD distillation from the AR checkpoint. + +### T2V Training + +```bash +torchrun --standalone --nnodes=1 --nproc_per_node=8 train.py \ + --config_path configs/train_ar.yaml \ + --logdir logs/train_ar \ + --wandb-save-dir wandb \ + --disable-wandb + +torchrun --standalone --nnodes=1 --nproc_per_node=8 train.py \ + --config_path configs/train_dmd.yaml \ + --logdir logs/train_dmd \ + --wandb-save-dir wandb \ + --disable-wandb +``` + +### I2V Training + +```bash +torchrun --standalone --nnodes=1 --nproc_per_node=8 train.py \ + --config_path configs/train_i2v_ar.yaml \ + --logdir logs/train_i2v_ar \ + --wandb-save-dir wandb \ + --disable-wandb + +torchrun --standalone --nnodes=1 --nproc_per_node=8 train.py \ + --config_path configs/train_i2v_dmd.yaml \ + --logdir logs/train_i2v_dmd \ + --wandb-save-dir wandb \ + --disable-wandb +``` + +For I2V configs, set `algorithm.i2v: true` and `algorithm.independent_first_frame: true`. `data.image_or_video_shape[1]` is the full latent sequence length, for example `96`, not `96 + 1`: the clean image latent replaces the first latent during denoising and that first latent is masked out of the training loss. For I2V DMD, set `checkpoints.generator_ckpt` to the I2V AR checkpoint used to initialize the student. + +## Models + +| Model | FPS ↑ | Params | VBench ↑ | Multi-shot | +| --- | ---: | ---: | ---: | :---: | +| [LongLive-1.3B](https://huggingface.co/Efficient-Large-Model/LongLive-1.3B) | 20.7 | 1.3B | 84.87 | | +| [LongLive-2.0-5B](https://huggingface.co/Efficient-Large-Model/LongLive-2.0-5B) | 24.8 | 5B | 85.06 | ✅ | +| [LongLive-2.0-5B-NVFP4-4Step](https://huggingface.co/Efficient-Large-Model/LongLive-2.0-5B-NVFP4-S4) | 29.7 | 5B | 84.51 | ✅ | +| [LongLive-2.0-5B-NVFP4-2Step](https://huggingface.co/Efficient-Large-Model/LongLive-2.0-5B-NVFP4-S2) | 45.7 | 5B | 83.14 | ✅ | + +## Awesome work using LongLive + +- [DreamForge-World 0.1](https://trydreamforge.com/): Adapts the LongLive AR video stack with a residual action pathway for low-compute real-time controllable world modeling. +- [DreamX-World 1.0](https://arxiv.org/abs/2606.16993): Follows LongLive by adapting the model on long sequences with long rollouts and local temporal windows for stable long-horizon AR world generation. +- [SANA-Video](https://nvlabs.github.io/Sana/docs/longsana/): Combines SANA-Video with LongLive to build LongSANA, a real-time minute-long video generation variant with constant-memory KV cache. +- [Daydream Scope](https://docs.daydream.live/scope/reference/pipelines/longlive): Wraps LongLive as a streaming AR video diffusion pipeline for interactive text-to-video and video-to-video workflows. +- [MemFlow](https://github.com/KlingAIResearch/MemFlow): Builds on the LongLive codebase and adds adaptive memory retrieval for more consistent long narrative video generation. +- [ShotStream](https://github.com/KlingAIResearch/ShotStream): Builds on LongLive’s distillation procedure for real-time streaming multi-shot AR video generation. +- [Stream-T1](https://github.com/FrameX-AI/Stream-T1): Builds on LongLive's codebase and algorithm, adding test-time scaling with noise propagation, reward pruning, and memory sinking. +- [KVPO](https://github.com/Richard-Zhang-AI/KVPO): Builds on LongLive and related AR video codebases to perform GRPO-style alignment through historical KV semantic exploration. +- [LoL](https://github.com/justincui03/LoL): Builds on LongLive to study and mitigate sink-collapse for ultra-long AR streaming video generation. +- [TriAttention](https://github.com/WeianMao/triattention/tree/main/longlive): Integrates trigonometric KV-cache compression into LongLive's causal inference pipeline, reducing KV memory inside LongLive's local-attention window. +- [StreamEdit](https://github.com/DSL-Lab/StreamEdit): Provides a `LongLive_StreamEdit` implementation for training-free streaming video editing built on the LongLive v1.0 codebase. +- [Streaming Autoregressive Video Generation via Diagonal Distillation](https://github.com/Sphere-AI-Lab/diagdistill): Builds on the LongLive codebase and supports direct initialization from `LongLive-1.3B` checkpoints for streaming AR video distillation. +- [Forcing-KV](https://github.com/zju-jiyicheng/Forcing-KV): Adds hybrid KV-cache compression to LongLive, including LongLive inference and interactive-generation scripts. +- [Dummy Forcing](https://github.com/csguoh/DummyForcing): Unifies Self-Forcing, LongLive, and Causal-Forcing pipelines with LongLive inference, VBench, and interactive-generation configs. +- [MemRoPE](https://github.com/YoungRaeKimm/MemRoPE): Uses LongLive as a supported base model for training-free infinite video generation with evolving memory tokens. +- [Astrolabe](https://github.com/franklinz233/Astrolabe): Supports LongLive as a distilled autoregressive video backbone with LongLive-specific RL configs and LoRA initialization. + + +## License +This repository is released under the Apache 2.0 license. See [LICENSE](LICENSE) for details. + +## Citation +Please consider citing our work if you find them useful: + +```bibtex +@article{longlive_2.0, + title={LongLive2.0: An NVFP4 Parallel Infrastructure for Long Video Generation}, + author={Chen, Yukang and Wang, Luozhou and Huang, Wei and Yang, Shuai and Zhang, Bohan and Xiao, Yicheng and Chu, Ruihang and Mao, Weian and Hu, Qixin and Liu, Shaoteng and Zhao, Yuyang and Mao, Huizi and Chen, Ying-Cong and Xie, Enze and Qi, Xiaojuan and Han, Song}, + journal={arXiv preprint arXiv: 2605.18739}, + year={2026} +} +``` + +```bibtex +@inproceedings{longlive, + title={Longlive: Real-time interactive long video generation}, + author={Yang, Shuai and Huang, Wei and Chu, Ruihang and Xiao, Yicheng and Zhao, Yuyang and Wang, Xianbang and Li, Muyang and Xie, Enze and Chen, Yingcong and Lu, Yao and others}, + booktitle={ICLR}, + year={2026}, +} +``` + +```bibtex +@article{longlive_rag, + title = {LongLive-RAG: A General Retrieval-Augmented Framework for Long Video Generation}, + author = {Hu, Qixin and Yang, Shuai and Huang, Wei and Han, Song and Chen, Yukang}, + journal = {arXiv preprint arXiv:2606.02553}, + year = {2026} +} +``` + +## Acknowledgement +- [Self-Forcing](https://github.com/guandeh17/Self-Forcing): the AR training codebase and formulation we build upon. +- [Wan2.2](https://github.com/Wan-Video/Wan2.2): the base video diffusion model components used in this release. diff --git a/README.wehub.md b/README.wehub.md new file mode 100644 index 0000000..b66babe --- /dev/null +++ b/README.wehub.md @@ -0,0 +1,7 @@ +# WeHub 来源说明 + +- 原始项目：`NVlabs/LongLive` +- 原始仓库：https://github.com/NVlabs/LongLive +- 导入方式：上游默认分支的最新快照 +- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准 +- 本文件仅用于记录来源，不代表 WeHub 是原项目作者 diff --git a/assets/longlive2/LongLive1_teaser.png b/assets/longlive2/LongLive1_teaser.png new file mode 100644 index 0000000..af5a346 Binary files /dev/null and b/assets/longlive2/LongLive1_teaser.png differ diff --git a/assets/longlive2/fig_framework_overview.png b/assets/longlive2/fig_framework_overview.png new file mode 100644 index 0000000..1c7509d Binary files /dev/null and b/assets/longlive2/fig_framework_overview.png differ diff --git a/assets/longlive2/first-video-frame.png b/assets/longlive2/first-video-frame.png new file mode 100644 index 0000000..394463d Binary files /dev/null and b/assets/longlive2/first-video-frame.png differ diff --git a/assets/longlive2/logo.png b/assets/longlive2/logo.png new file mode 100644 index 0000000..a54631d Binary files /dev/null and b/assets/longlive2/logo.png differ diff --git a/assets/longlive2/teaser.jpg b/assets/longlive2/teaser.jpg new file mode 100644 index 0000000..7930ad4 Binary files /dev/null and b/assets/longlive2/teaser.jpg differ diff --git a/configs/fp8/inference_fp8.yaml b/configs/fp8/inference_fp8.yaml new file mode 100644 index 0000000..bd96545 --- /dev/null +++ b/configs/fp8/inference_fp8.yaml @@ -0,0 +1,51 @@ +# TorchAO FP8 PTQ inference from the merged BF16 LongLive checkpoint. +model_kwargs: + model_name: Wan2.2-TI2V-5B + timestep_shift: 5.0 + num_frame_per_block: 8 + local_attn_size: 32 + +use_ema: false +output_folder: videos/longlive2_fp8 +num_samples: 1 +num_output_frames: 8 +save_latents_only: false +inference_iter: -1 + +data: + data_path: example/long_example.txt + image_or_video_shape: + - 1 + - 8 + - 48 + - 44 + - 80 + +inference: + sampling_steps: 4 + sink_size: 8 + guidance_scale: 1.0 + multi_shot_sink: true + multi_shot_rope_offset: 8 + streaming_vae: false + async_vae: false + vae_type: wan + +checkpoints: + generator_ckpt: /path/to/LongLive-2.0-5B/model_bf16.pt + +# Load the checkpoint in BF16, then quantize compatible generator Linear layers +# to row-wise dynamic FP8 (W8A8) before torch.compile. +fp8_quant: true + +torch_compile: auto +torch_compile_min_samples: 3 +torch_compile_target: generator_model +torch_compile_backend: inductor +torch_compile_mode: max-autotune-no-cudagraphs +torch_compile_fullgraph: false +torch_compile_dynamic: false +torch_compile_suppress_errors: false + +logging: + seed: 0 diff --git a/configs/inference.yaml b/configs/inference.yaml new file mode 100644 index 0000000..3d00676 --- /dev/null +++ b/configs/inference.yaml @@ -0,0 +1,48 @@ +# Passed directly to WanDiffusionWrapper(**model_kwargs). +model_kwargs: + model_name: Wan2.2-TI2V-5B + timestep_shift: 5.0 + num_frame_per_block: 8 + local_attn_size: 32 + +use_ema: false +output_folder: videos/longlive2 +num_samples: 1 +save_latents_only: false + +data: + data_path: /path/to/longlive2/inference_prompts + # Latent tensor shape [B, F, C, H, W]. F is also the default number of + # generated latent frames unless num_output_frames is provided. + image_or_video_shape: + - 1 + - 128 + - 48 + - 44 + - 80 + +inference: + sampling_steps: 4 + sink_size: 8 + guidance_scale: 1.0 + multi_shot_sink: true + multi_shot_rope_offset: 8 + streaming_vae: false + async_vae: false + vae_type: wan # wan, mg_lightvae, mg_lightvae_v2 + vae_device: "cuda:2" + +checkpoints: + generator_ckpt: /path/to/longlive2/generator_checkpoint.pt + lora_ckpt: /path/to/longlive2/lora_checkpoint.pt + +# Presence of this section enables LoRA inference; remove it for full-generator inference. +adapter: + type: lora + rank: 128 + alpha: 128 + dropout: 0.0 + verbose: true + +logging: + seed: 0 diff --git a/configs/inference_sp.yaml b/configs/inference_sp.yaml new file mode 100644 index 0000000..e755db6 --- /dev/null +++ b/configs/inference_sp.yaml @@ -0,0 +1,51 @@ +# Ulysses sequence-parallel inference config for Wan2.2-TI2V-5B. +# +# Example launch: +# torchrun --nproc_per_node=4 inference_sp.py --config_path configs/inference_sp.yaml + +model_kwargs: + model_name: Wan2.2-TI2V-5B + timestep_shift: 5.0 + num_frame_per_block: 8 + local_attn_size: 32 + +sp_size: 4 +dp_size: 1 +auto_sp_remainder: false +model_num_heads: 24 + +use_ema: false +output_folder: videos/longlive2_sp +num_samples: 1 +save_latents_only: false + +data: + data_path: /path/to/longlive2/inference_prompts + image_or_video_shape: + - 1 + - 128 + - 48 + - 44 + - 80 + +inference: + sampling_steps: 4 + sink_size: 8 + guidance_scale: 1.0 + multi_shot_sink: true + multi_shot_rope_offset: 8 + +checkpoints: + generator_ckpt: /path/to/longlive2/generator_checkpoint.pt + lora_ckpt: /path/to/longlive2/lora_checkpoint.pt + +# Presence of this section enables LoRA inference; remove it for full-generator inference. +adapter: + type: lora + rank: 128 + alpha: 128 + dropout: 0.0 + verbose: true + +logging: + seed: 0 diff --git a/configs/nvfp4/inference_i2v_nvfp4.yaml b/configs/nvfp4/inference_i2v_nvfp4.yaml new file mode 100644 index 0000000..b31ba56 --- /dev/null +++ b/configs/nvfp4/inference_i2v_nvfp4.yaml @@ -0,0 +1,94 @@ +# NVFP4 i2v inference with optional KV cache quantization and streaming VAE. +# = configs/nvfp4/inference_nvfp4.yaml + i2v conditioning. +# The first latent frame is encoded from the source image (provided by the +# video dataset), kept clean during denoising, and the rollout continues from it. +model_kwargs: + model_name: Wan2.2-TI2V-5B + timestep_shift: 5.0 + num_frame_per_block: 8 + local_attn_size: 32 + sink_size: 8 + +# Enable i2v: build the video dataset, encode the first frame to a clean latent, +# and clamp it as the first-chunk conditioning during inference. +i2v: true + +use_ema: false +output_folder: videos/longlive2_i2v_nvfp4 +num_samples: 1 +save_latents_only: false +save_with_index: true +inference_iter: -1 +num_output_frames: 384 +merge_lora: false + +# i2v dataset / first-frame handling (read by MultiVideoConcatDataset in inference.py). +allow_padding: false +min_latent_frames: 0 +max_chunks_per_shot: 0 +uniform_prompt: false + +data: + # For i2v the data_path points at a video/image dataset; the first frame of + # each clip is used as the conditioning image. + data_path: /path/to/longlive2/inference_video_dataset + image_or_video_shape: + - 1 + - 384 + - 48 + - 44 + - 80 + +inference: + sampling_steps: 4 + # i2v conditioning: keep the first latent frame independent and clean. + independent_first_frame: true + # i2v training-time eval uses sink_size: 0; sink_size: 8 here matches the + # long-context streaming NVFP4 t2v setup. Tune as needed. + sink_size: 8 + guidance_scale: 1.0 + multi_shot_sink: true + multi_shot_rope_offset: 8 + kv_quant: true + kv_quant_scale_rule: mse + kv_quant_backend: cuda + streaming_vae: false + async_vae: false + vae_type: wan # wan, mg_lightvae, mg_lightvae_v2 + vae_device: "cuda:2" + negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走' + +checkpoints: + generator_ckpt: /path/to/model_te.pt + # generator_ckpt: /path/to/model_4o6.pt # only set this if you are using FourOverSix (ckpt) for NVFP4 + +model_quant: true +model_quant_use_transformer_engine: true # only set this to true if you are using Transformer Engine (ckpt) for NVFP4; otherwise, set it to false (4o6) +model_quant_te_inference_only: true +model_quant_te_low_precision_weights: true +model_quant_te_fallback_to_fouroversix: false # only set this to false if you are using FourOverSix (ckpt) for NVFP4; otherwise, set it to true (TE) +model_quant_scale_rule: mse +model_quant_activation_scale_rule: mse +model_quant_weight_scale_rule: mse +model_quant_gradient_scale_rule: mse + +torch_compile: auto +torch_compile_min_samples: 2 +torch_compile_target: generator_model +torch_compile_backend: inductor +torch_compile_mode: max-autotune-no-cudagraphs +torch_compile_fullgraph: false +torch_compile_dynamic: false +torch_compile_suppress_errors: true + +adapter: + type: lora + rank: 128 + alpha: 128 + dropout: 0.0 + dtype: bfloat16 + apply_to_critic: true + verbose: true + +logging: + seed: 0 diff --git a/configs/nvfp4/inference_nvfp4.yaml b/configs/nvfp4/inference_nvfp4.yaml new file mode 100644 index 0000000..c5420a5 --- /dev/null +++ b/configs/nvfp4/inference_nvfp4.yaml @@ -0,0 +1,75 @@ +# NVFP4 inference with optional KV cache quantization and streaming VAE. +model_kwargs: + model_name: Wan2.2-TI2V-5B + timestep_shift: 5.0 + num_frame_per_block: 8 + local_attn_size: 32 + sink_size: 8 + +use_ema: false +output_folder: videos/longlive2_nvfp4 +num_samples: 1 +save_latents_only: false +save_with_index: true +inference_iter: -1 +num_output_frames: 384 +merge_lora: false + +data: + data_path: /path/to/longlive2/inference_prompts + image_or_video_shape: + - 1 + - 384 + - 48 + - 44 + - 80 + +inference: + sampling_steps: 4 + sink_size: 8 + guidance_scale: 1.0 + multi_shot_sink: true + multi_shot_rope_offset: 8 + kv_quant: true + kv_quant_scale_rule: mse + kv_quant_backend: cuda + streaming_vae: false + async_vae: false + vae_type: wan # wan, mg_lightvae, mg_lightvae_v2 + vae_device: "cuda:2" + negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走' + +checkpoints: + generator_ckpt: /path/to/model_te.pt + # generator_ckpt: /path/to/model_4o6.pt # only set this if you are using FourOverSix (ckpt) for NVFP4 + +model_quant: true +model_quant_use_transformer_engine: true # only set this to true if you are using Transformer Engine (ckpt) for NVFP4; otherwise, set it to false (4o6) +model_quant_te_inference_only: true +model_quant_te_low_precision_weights: true +model_quant_te_fallback_to_fouroversix: false # only set this to false if you are using FourOverSix (ckpt) for NVFP4; otherwise, set it to true (TE) +model_quant_scale_rule: mse +model_quant_activation_scale_rule: mse +model_quant_weight_scale_rule: mse +model_quant_gradient_scale_rule: mse + +torch_compile: auto +torch_compile_min_samples: 2 +torch_compile_target: generator_model +torch_compile_backend: inductor +torch_compile_mode: max-autotune-no-cudagraphs +torch_compile_fullgraph: false +torch_compile_dynamic: false +torch_compile_suppress_errors: true + +adapter: + type: lora + rank: 128 + alpha: 128 + dropout: 0.0 + dtype: bfloat16 + apply_to_critic: true + verbose: true + +logging: + seed: 0 diff --git a/configs/nvfp4/train_ar_nvfp4.yaml b/configs/nvfp4/train_ar_nvfp4.yaml new file mode 100644 index 0000000..9019c0c --- /dev/null +++ b/configs/nvfp4/train_ar_nvfp4.yaml @@ -0,0 +1,96 @@ +# NVFP4 AR teacher-forcing training. +infra: + sequence_parallel_size: 4 + sharding_strategy: hybrid_full + mixed_precision: true + gradient_checkpointing: true + vae_halo_latents: 28 + generator_fsdp_wrap_strategy: size + text_encoder_fsdp_wrap_strategy: size + model_quant: true + model_quant_scale_rule: static_6 + model_quant_activation_scale_rule: mse + model_quant_weight_scale_rule: static_6 + model_quant_gradient_scale_rule: mse + +model_kwargs: + model_name: Wan2.2-TI2V-5B + timestep_shift: 5.0 + num_frame_per_block: 8 + local_attn_size: -1 + sink_size: 8 + t_scale: 1.0 + rope_method: linear + +checkpoints: + generator_ckpt: /path/to/wan_or_ar_checkpoint.pt + +algorithm: + trainer: diffusion + causal: true + teacher_forcing: true + num_train_timestep: 1000 + denoising_loss_type: flow + +training: + lr: 1.0e-05 + weight_decay: 0.0 + beta1: 0.0 + beta2: 0.999 + batch_size: 1 + gradient_accumulation_steps: 4 + ema_weight: 0.99 + ema_start_step: 100 + log_iters: 15 + max_checkpoints: 10 + max_iters: 600 + gc_interval: 100 + error_recycling: + enabled: true + num_buckets: 50 + buffer_size_per_bucket: 500 + buffer_warmup_iter: 50 + context_inject_prob: 0.9 + latent_inject_prob: 0.9 + noise_inject_prob: 0.01 + clean_prob: 0.2 + clean_buffer_update_prob: 0.1 + modulate_factor: 0.3 + start_step: 0 + replacement_strategy: random + +data: + data_path: /path/to/longlive2/train_dataset_64s + eval_data_path: /path/to/longlive2/eval_prompts + image_or_video_shape: + - 1 + - 384 + - 48 + - 44 + - 80 + load_raw_video: true + max_chunks_per_shot: 4 + +inference: + sampling_steps: 50 + sink_size: 8 + local_attn_size: 24 + multi_shot_sink: true + multi_shot_rope_offset: 8 + guidance_scale: 3.0 + use_relative_rope: false + inference_t_scale: 1.0 + negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走' + +evaluation: + interval: 1000 + num_frames: 384 + use_ema: true + val_batch_size: 1 + save_latents_only: true + +logging: + seed: 0 + wandb_key: null + wandb_entity: null + wandb_project: LongLive2-AR-NVFP4 diff --git a/configs/nvfp4/train_dmd_nvfp4_step4.yaml b/configs/nvfp4/train_dmd_nvfp4_step4.yaml new file mode 100644 index 0000000..d28cfed --- /dev/null +++ b/configs/nvfp4/train_dmd_nvfp4_step4.yaml @@ -0,0 +1,126 @@ +# NVFP4 DMD LoRA distillation, 4-step rollout. +infra: + sharding_strategy: hybrid_full + mixed_precision: true + gradient_checkpointing: true + generator_fsdp_wrap_strategy: size + real_score_fsdp_wrap_strategy: size + fake_score_fsdp_wrap_strategy: size + text_encoder_fsdp_wrap_strategy: size + generator_quant: true + real_score_quant: true + fake_score_quant: false + real_score_quant_materialize: true + generator_quant_scale_rule: mse + generator_quant_activation_scale_rule: mse + generator_quant_weight_scale_rule: mse + generator_quant_gradient_scale_rule: mse + generator_quant_use_default_filtered_modules: false + real_score_quant_scale_rule: mse + real_score_quant_activation_scale_rule: mse + real_score_quant_weight_scale_rule: mse + real_score_quant_use_default_filtered_modules: false + generator_quant_filtered_modules: &nvfp4_filtered_modules + - text_embedding.0 + - text_embedding.2 + - patch_embedding + - time_projection.1 + - time_embedding.0 + - time_embedding.2 + - head.head + - head.modulation + - re:.*norm_k$ + - re:.*norm_q$ + - re:.*norm1$ + - re:.*norm2$ + - re:.*norm3$ + real_score_quant_filtered_modules: *nvfp4_filtered_modules + +model_kwargs: + model_name: Wan2.2-TI2V-5B + timestep_shift: 5.0 + num_frame_per_block: 8 + local_attn_size: 32 + sink_size: 8 + +checkpoints: + generator_ckpt: null + real_score_ckpt: null + fake_score_ckpt: null + +algorithm: + trainer: score_distillation + distribution_loss: dmd + all_causal: true + ts_schedule: false + num_train_timestep: 1000 + denoising_loss_type: flow + warp_denoising_step: false + denoising_step_list: [1000, 946, 854, 681, 0] + teacher_forcing: false + backward_simulation: true + real_guidance_scale: 3.0 + fake_guidance_scale: 0.0 + +training: + lr: 1.0e-05 + lr_critic: 2.0e-06 + weight_decay: 0.0 + beta1: 0.0 + beta2: 0.999 + beta1_critic: 0.0 + beta2_critic: 0.999 + batch_size: 2 + gradient_accumulation_steps: 1 + ema_weight: 0.99 + ema_start_step: 200 + log_iters: 50 + max_checkpoints: 20 + max_iters: 5000 + gc_interval: 100 + dfake_gen_update_ratio: 5 + min_num_training_frames: 32 + num_training_frames: 32 + slice_last_frames: 32 + chunks_per_shot: 2 + +data: + data_path: /path/to/longlive2/train_dataset_16to32s + eval_data_path: /path/to/longlive2/eval_prompts + image_or_video_shape: + - 1 + - 32 + - 48 + - 44 + - 80 + load_raw_video: true + +inference: + sampling_steps: 4 + sink_size: 8 + multi_shot_sink: true + multi_shot_rope_offset: 8 + narrative_rope_phase_shift: 8 + negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走' + +evaluation: + interval: 99999999999 + num_frames: 32 + use_ema: false + val_batch_size: 1 + save_latents_only: true + +adapter: + type: lora + rank: 128 + alpha: 128 + dropout: 0.0 + dtype: bfloat16 + apply_to_critic: true + verbose: true + +logging: + seed: 0 + wandb_key: null + wandb_entity: null + wandb_project: LongLive2-DMD-NVFP4 diff --git a/configs/nvfp4/train_i2v_ar_nvfp4.yaml b/configs/nvfp4/train_i2v_ar_nvfp4.yaml new file mode 100644 index 0000000..8864d36 --- /dev/null +++ b/configs/nvfp4/train_i2v_ar_nvfp4.yaml @@ -0,0 +1,101 @@ +# NVFP4 i2v AR teacher-forcing training. +# = configs/train_i2v_ar.yaml + the NVFP4 quant infra from +# configs/nvfp4/train_ar_nvfp4.yaml. +infra: + sequence_parallel_size: 4 + sharding_strategy: hybrid_full + mixed_precision: true + gradient_checkpointing: true + vae_halo_latents: 28 + generator_fsdp_wrap_strategy: size + text_encoder_fsdp_wrap_strategy: size + model_quant: true + model_quant_scale_rule: static_6 + model_quant_activation_scale_rule: mse + model_quant_weight_scale_rule: static_6 + model_quant_gradient_scale_rule: mse + +model_kwargs: + model_name: Wan2.2-TI2V-5B + timestep_shift: 5.0 + num_frame_per_block: 8 + local_attn_size: -1 + sink_size: 8 + t_scale: 1.0 + rope_method: linear + +checkpoints: + generator_ckpt: /path/to/wan_or_ar_checkpoint.pt + +# I2V AR diffusion objective. The first latent frame is encoded from the +# source image, kept clean, and masked out of the flow-matching loss. +algorithm: + trainer: diffusion + i2v: true + causal: true + teacher_forcing: true + independent_first_frame: true + num_train_timestep: 1000 + denoising_loss_type: flow + +training: + lr: 1.0e-05 + weight_decay: 0.0 + beta1: 0.0 + beta2: 0.999 + batch_size: 1 + gradient_accumulation_steps: 1 + ema_weight: 0.99 + ema_start_step: 100 + log_iters: 15 + max_checkpoints: 20 + max_iters: 600 + gc_interval: 100 + error_recycling: + enabled: true + num_buckets: 50 + buffer_size_per_bucket: 500 + buffer_warmup_iter: 50 + context_inject_prob: 0.9 + latent_inject_prob: 0.9 + noise_inject_prob: 0.01 + clean_prob: 0.2 + clean_buffer_update_prob: 0.1 + modulate_factor: 0.3 + start_step: 0 + replacement_strategy: random + +data: + data_path: /path/to/longlive2/train_video_dataset + eval_data_path: /path/to/longlive2/eval_video_dataset + image_or_video_shape: + - 1 + - 96 + - 48 + - 44 + - 80 + load_raw_video: true + +inference: + sampling_steps: 50 + sink_size: 0 + multi_shot_sink: true + multi_shot_rope_offset: 8 + guidance_scale: 3.0 + independent_first_frame: true + use_relative_rope: false + inference_t_scale: 1.0 + negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走' + +evaluation: + interval: 1000 + num_frames: 32 + use_ema: true + val_batch_size: 1 + save_latents_only: true + +logging: + seed: 0 + wandb_key: null + wandb_entity: null + wandb_project: LongLive2-I2V-AR-NVFP4 diff --git a/configs/nvfp4/train_i2v_dmd_nvfp4_step4.yaml b/configs/nvfp4/train_i2v_dmd_nvfp4_step4.yaml new file mode 100644 index 0000000..da01470 --- /dev/null +++ b/configs/nvfp4/train_i2v_dmd_nvfp4_step4.yaml @@ -0,0 +1,142 @@ +# NVFP4 i2v DMD LoRA distillation, 4-step rollout. +# = configs/train_i2v_dmd.yaml + the NVFP4 quant infra and 4-step denoising +# schedule from configs/nvfp4/train_dmd_nvfp4_step4.yaml. +infra: + sharding_strategy: hybrid_full + mixed_precision: true + gradient_checkpointing: true + generator_fsdp_wrap_strategy: size + real_score_fsdp_wrap_strategy: size + fake_score_fsdp_wrap_strategy: size + text_encoder_fsdp_wrap_strategy: size + generator_quant: true + real_score_quant: true + fake_score_quant: false + real_score_quant_materialize: true + generator_quant_scale_rule: mse + generator_quant_activation_scale_rule: mse + generator_quant_weight_scale_rule: mse + generator_quant_gradient_scale_rule: mse + generator_quant_use_default_filtered_modules: false + real_score_quant_scale_rule: mse + real_score_quant_activation_scale_rule: mse + real_score_quant_weight_scale_rule: mse + real_score_quant_use_default_filtered_modules: false + generator_quant_filtered_modules: &nvfp4_filtered_modules + - text_embedding.0 + - text_embedding.2 + - patch_embedding + - time_projection.1 + - time_embedding.0 + - time_embedding.2 + - head.head + - head.modulation + - re:.*norm_k$ + - re:.*norm_q$ + - re:.*norm1$ + - re:.*norm2$ + - re:.*norm3$ + real_score_quant_filtered_modules: *nvfp4_filtered_modules + +model_kwargs: + model_name: Wan2.2-TI2V-5B + timestep_shift: 5.0 + num_frame_per_block: 8 + local_attn_size: 32 + sink_size: 8 + +# Optional initialization checkpoints. In practice, set generator_ckpt to the +# i2v AR checkpoint used as the DMD student initialization. +checkpoints: + generator_ckpt: null + real_score_ckpt: null + fake_score_ckpt: null + +# I2V DMD objective. Backward simulation uses the same first-chunk image +# conditioning rule as i2v inference: the clean first latent is injected during +# denoising and excluded from generator/critic losses. +algorithm: + trainer: score_distillation + distribution_loss: dmd + i2v: true + all_causal: true + ts_schedule: false + num_train_timestep: 1000 + denoising_loss_type: flow + warp_denoising_step: false + denoising_step_list: [1000, 946, 854, 681, 0] + teacher_forcing: false + backward_simulation: true + independent_first_frame: true + real_guidance_scale: 3.0 + fake_guidance_scale: 0.0 + +training: + lr: 1.0e-05 + lr_critic: 2.0e-06 + weight_decay: 0.0 + beta1: 0.0 + beta2: 0.999 + beta1_critic: 0.0 + beta2_critic: 0.999 + batch_size: 1 + gradient_accumulation_steps: 1 + ema_weight: 0.99 + ema_start_step: 200 + log_iters: 50 + max_checkpoints: 50 + max_iters: 5000 + gc_interval: 100 + dfake_gen_update_ratio: 5 + # Backward simulation samples a rollout length in [min_num_training_frames, num_training_frames]. + # Setting both to 32 gives a fixed 32-latent DMD rollout. + min_num_training_frames: 32 + num_training_frames: 32 + # Keep this many tail latents for the DMD/critic losses after backward simulation. + # With a 32-latent rollout, 32 keeps the full generated window. + slice_last_frames: 32 + chunks_per_shot: 2 + +data: + data_path: /path/to/longlive2/train_video_dataset + eval_data_path: /path/to/longlive2/eval_video_dataset + image_or_video_shape: + - 1 + - 32 + - 48 + - 44 + - 80 + load_raw_video: true + +# The distillation loop generates i2v trajectories, then computes DMD losses. +inference: + sampling_steps: 4 + sink_size: 0 + multi_shot_sink: true + multi_shot_rope_offset: 8 + independent_first_frame: true + negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走' + +# Optional validation generation from the distilled model. Uses inference sampling/cache settings by default. +evaluation: + interval: 99999999999 + num_frames: 32 + use_ema: false + val_batch_size: 1 + save_latents_only: true + +# Presence of this section enables LoRA distillation; remove it for full fine-tuning. +adapter: + type: lora + rank: 128 + alpha: 128 + dropout: 0.0 + dtype: bfloat16 + apply_to_critic: true + verbose: true + +logging: + seed: 0 + wandb_key: null + wandb_entity: null + wandb_project: LongLive2-I2V-DMD-NVFP4 diff --git a/configs/train_ar.yaml b/configs/train_ar.yaml new file mode 100644 index 0000000..51a3ffc --- /dev/null +++ b/configs/train_ar.yaml @@ -0,0 +1,85 @@ +# Distributed training runtime. +infra: + sequence_parallel_size: 4 + sharding_strategy: hybrid_full + mixed_precision: true + gradient_checkpointing: true + vae_halo_latents: 28 + generator_fsdp_wrap_strategy: size + text_encoder_fsdp_wrap_strategy: size + +# Passed directly to WanDiffusionWrapper(**model_kwargs). +model_kwargs: + model_name: Wan2.2-TI2V-5B + timestep_shift: 5.0 + num_frame_per_block: 8 + local_attn_size: -1 + +# AR diffusion objective. +algorithm: + trainer: diffusion + causal: true + teacher_forcing: true + num_train_timestep: 1000 + denoising_loss_type: flow + +training: + lr: 1.0e-05 + weight_decay: 0.0 + beta1: 0.0 + beta2: 0.999 + batch_size: 1 + gradient_accumulation_steps: 1 + ema_weight: 0.99 + ema_start_step: 100 + log_iters: 1 + max_checkpoints: 20 + max_iters: 600 + gc_interval: 100 + error_recycling: + enabled: true + num_buckets: 50 + buffer_size_per_bucket: 500 + buffer_warmup_iter: 50 + context_inject_prob: 0.9 + latent_inject_prob: 0.9 + noise_inject_prob: 0.01 + clean_prob: 0.2 + clean_buffer_update_prob: 0.1 + modulate_factor: 0.3 + start_step: 0 + replacement_strategy: random + +data: + data_path: /path/to/longlive2/train_dataset + eval_data_path: /path/to/longlive2/eval_prompts + image_or_video_shape: + - 1 + - 96 + - 48 + - 44 + - 80 + load_raw_video: true + +# Generation settings used by training-time evaluation. +inference: + sampling_steps: 50 + sink_size: 0 + multi_shot_sink: true + multi_shot_rope_offset: 8 + guidance_scale: 3.0 + negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走' + +# Lightweight validation generation during AR training. +evaluation: + interval: 1 + num_frames: 16 + use_ema: true + val_batch_size: 1 + save_latents_only: true + +logging: + seed: 0 + wandb_key: null + wandb_entity: null + wandb_project: LongLive2-AR diff --git a/configs/train_dmd.yaml b/configs/train_dmd.yaml new file mode 100644 index 0000000..5ec001b --- /dev/null +++ b/configs/train_dmd.yaml @@ -0,0 +1,102 @@ +# Distributed distillation runtime. +infra: + sharding_strategy: hybrid_full + mixed_precision: true + gradient_checkpointing: true + generator_fsdp_wrap_strategy: size + real_score_fsdp_wrap_strategy: size + fake_score_fsdp_wrap_strategy: size + text_encoder_fsdp_wrap_strategy: size + +# Shared Wan backbone settings for the student, real-score teacher, and fake-score critic. +# Role-specific model kwargs can still be provided as real_model_kwargs/fake_model_kwargs +# if a future run needs different model construction. +model_kwargs: + model_name: Wan2.2-TI2V-5B + timestep_shift: 5.0 + num_frame_per_block: 8 + local_attn_size: 32 + +# Optional initialization checkpoints. These are flattened by normalize_config +# so the trainer still reads generator_ckpt/real_score_ckpt/fake_score_ckpt. +# With adapter enabled, generator_ckpt initializes the base generator before LoRA wrapping. +checkpoints: + generator_ckpt: null + real_score_ckpt: null + fake_score_ckpt: null + +# DMD distillation objective. +algorithm: + trainer: score_distillation + distribution_loss: dmd + all_causal: true + ts_schedule: false + real_guidance_scale: 3.0 + fake_guidance_scale: 0.0 + + +training: + lr: 1.0e-05 + lr_critic: 2.0e-06 + weight_decay: 0.0 + beta1: 0.0 + beta2: 0.999 + beta1_critic: 0.0 + beta2_critic: 0.999 + batch_size: 1 + gradient_accumulation_steps: 1 + ema_weight: 0.99 + ema_start_step: 200 + log_iters: 50 + max_checkpoints: 50 + max_iters: 5000 + gc_interval: 100 + + dfake_gen_update_ratio: 5 + # Backward simulation samples a rollout length in [min_num_training_frames, num_training_frames]. + # Setting both to 32 gives a fixed 32-latent DMD rollout. + min_num_training_frames: 32 + num_training_frames: 32 + # Keep this many tail latents for the DMD/critic losses after backward simulation. + # With a 32-latent rollout, 32 keeps the full generated window. + slice_last_frames: 32 + +data: + data_path: /path/to/longlive2/train_dataset + eval_data_path: /path/to/longlive2/eval_prompts + image_or_video_shape: + - 1 + - 32 + - 48 + - 44 + - 80 + +# The distillation loop generates video trajectories, then computes DMD losses. +inference: + sampling_steps: 4 + sink_size: 0 + multi_shot_rope_offset: 8 + +# Optional validation generation from the distilled model. Uses inference sampling/cache settings by default. +evaluation: + interval: 1 + num_frames: 32 + use_ema: false + # Can be raised when validation memory allows; 1 is the safe default. + val_batch_size: 1 + save_latents_only: true + +# Presence of this section enables LoRA distillation; remove it for full fine-tuning. +adapter: + type: lora + rank: 128 + alpha: 128 + dropout: 0.0 + apply_to_critic: true + verbose: true + +logging: + seed: 0 + wandb_key: null + wandb_entity: null + wandb_project: LongLive2-DMD diff --git a/configs/train_i2v_ar.yaml b/configs/train_i2v_ar.yaml new file mode 100644 index 0000000..3f88161 --- /dev/null +++ b/configs/train_i2v_ar.yaml @@ -0,0 +1,89 @@ +# Distributed i2v AR training runtime. +infra: + sequence_parallel_size: 4 + sharding_strategy: hybrid_full + mixed_precision: true + gradient_checkpointing: true + vae_halo_latents: 28 + generator_fsdp_wrap_strategy: size + text_encoder_fsdp_wrap_strategy: size + +# Passed directly to WanDiffusionWrapper(**model_kwargs). +model_kwargs: + model_name: Wan2.2-TI2V-5B + timestep_shift: 5.0 + num_frame_per_block: 8 + local_attn_size: -1 + +# I2V AR diffusion objective. The first latent frame is encoded from the +# source image, kept clean, and masked out of the flow-matching loss. +algorithm: + trainer: diffusion + i2v: true + causal: true + teacher_forcing: true + independent_first_frame: true + num_train_timestep: 1000 + denoising_loss_type: flow + +training: + lr: 1.0e-05 + weight_decay: 0.0 + beta1: 0.0 + beta2: 0.999 + batch_size: 1 + gradient_accumulation_steps: 1 + ema_weight: 0.99 + ema_start_step: 100 + log_iters: 1 + max_checkpoints: 20 + max_iters: 600 + gc_interval: 100 + error_recycling: + enabled: true + num_buckets: 50 + buffer_size_per_bucket: 500 + buffer_warmup_iter: 50 + context_inject_prob: 0.9 + latent_inject_prob: 0.9 + noise_inject_prob: 0.01 + clean_prob: 0.2 + clean_buffer_update_prob: 0.1 + modulate_factor: 0.3 + start_step: 0 + replacement_strategy: random + +data: + data_path: /path/to/longlive2/train_video_dataset + eval_data_path: /path/to/longlive2/eval_video_dataset + image_or_video_shape: + - 1 + - 96 + - 48 + - 44 + - 80 + load_raw_video: true + +# Generation settings used by training-time evaluation. +inference: + sampling_steps: 50 + sink_size: 0 + multi_shot_sink: true + multi_shot_rope_offset: 8 + guidance_scale: 3.0 + independent_first_frame: true + negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走' + +# Lightweight validation generation during i2v AR training. +evaluation: + interval: 1 + num_frames: 32 + use_ema: true + val_batch_size: 1 + save_latents_only: true + +logging: + seed: 0 + wandb_key: null + wandb_entity: null + wandb_project: LongLive2-I2V-AR diff --git a/configs/train_i2v_dmd.yaml b/configs/train_i2v_dmd.yaml new file mode 100644 index 0000000..e7df4d7 --- /dev/null +++ b/configs/train_i2v_dmd.yaml @@ -0,0 +1,104 @@ +# Distributed i2v DMD distillation runtime. +infra: + sharding_strategy: hybrid_full + mixed_precision: true + gradient_checkpointing: true + generator_fsdp_wrap_strategy: size + real_score_fsdp_wrap_strategy: size + fake_score_fsdp_wrap_strategy: size + text_encoder_fsdp_wrap_strategy: size + +# Shared Wan backbone settings for the student, real-score teacher, and fake-score critic. +model_kwargs: + model_name: Wan2.2-TI2V-5B + timestep_shift: 5.0 + num_frame_per_block: 8 + local_attn_size: 32 + +# Optional initialization checkpoints. In practice, set generator_ckpt to the +# i2v AR checkpoint used as the DMD student initialization. +checkpoints: + generator_ckpt: null + real_score_ckpt: null + fake_score_ckpt: null + +# I2V DMD objective. Backward simulation uses the same first-chunk image +# conditioning rule as i2v inference: the clean first latent is injected during +# denoising and excluded from generator/critic losses. +algorithm: + trainer: score_distillation + distribution_loss: dmd + i2v: true + all_causal: true + backward_simulation: true + teacher_forcing: false + independent_first_frame: true + ts_schedule: false + real_guidance_scale: 3.0 + fake_guidance_scale: 0.0 + +training: + lr: 1.0e-05 + lr_critic: 2.0e-06 + weight_decay: 0.0 + beta1: 0.0 + beta2: 0.999 + beta1_critic: 0.0 + beta2_critic: 0.999 + batch_size: 1 + gradient_accumulation_steps: 1 + ema_weight: 0.99 + ema_start_step: 200 + log_iters: 50 + max_checkpoints: 50 + max_iters: 5000 + gc_interval: 100 + dfake_gen_update_ratio: 5 + # Backward simulation samples a rollout length in [min_num_training_frames, num_training_frames]. + # Setting both to 32 gives a fixed 32-latent DMD rollout. + min_num_training_frames: 32 + num_training_frames: 32 + # Keep this many tail latents for the DMD/critic losses after backward simulation. + # With a 32-latent rollout, 32 keeps the full generated window. + slice_last_frames: 32 + +data: + data_path: /path/to/longlive2/train_video_dataset + eval_data_path: /path/to/longlive2/eval_video_dataset + image_or_video_shape: + - 1 + - 32 + - 48 + - 44 + - 80 + load_raw_video: true + +# The distillation loop generates i2v trajectories, then computes DMD losses. +inference: + sampling_steps: 4 + sink_size: 0 + multi_shot_rope_offset: 8 + independent_first_frame: true + +# Optional validation generation from the distilled model. Uses inference sampling/cache settings by default. +evaluation: + interval: 1 + num_frames: 32 + use_ema: false + val_batch_size: 1 + save_latents_only: true + +# Presence of this section enables LoRA distillation; remove it for full fine-tuning. +adapter: + type: lora + rank: 128 + alpha: 128 + dropout: 0.0 + apply_to_critic: true + verbose: true + +logging: + seed: 0 + wandb_key: null + wandb_entity: null + wandb_project: LongLive2-I2V-DMD diff --git a/docs/FLASH_ATTENTION_3_AND_HOPPER_SUPPORT.md b/docs/FLASH_ATTENTION_3_AND_HOPPER_SUPPORT.md new file mode 100644 index 0000000..c0e5903 --- /dev/null +++ b/docs/FLASH_ATTENTION_3_AND_HOPPER_SUPPORT.md @@ -0,0 +1,103 @@ +# Flash Attention 3 and Hopper GPU Support + +This document describes the Flash Attention 3 (FA3) integration and extended Hopper GPU support in LongLive. + +## Overview + +LongLive supports both Flash Attention 2 (FA2) and Flash Attention 3 (FA3) for efficient attention computation. FA3 is automatically enabled on Hopper architecture GPUs (Compute Capability 9.0+), providing improved performance. + +## Supported Hardware + +### Hopper Architecture GPUs (FA3 Enabled) +- **NVIDIA H100** - Data center GPU +- **NVIDIA H800** - China-specific variant +- **NVIDIA H20** - China-specific variant + +All Hopper GPUs share Compute Capability 9.0, which is the requirement for FA3. + +### Other GPUs (FA2 Fallback) +- **NVIDIA A100** - Ampere architecture (Compute Capability 8.0) +- **NVIDIA A800** - Ampere architecture (Compute Capability 8.0) +- Other CUDA-capable GPUs with FA2 support + +## Design Choices + +### 1. GPU Detection via Compute Capability + +Instead of relying on device name string matching (which would miss H800/H20), we detect Hopper GPUs using CUDA Compute Capability: + +```python +def is_hopper_gpu(): + if torch.cuda.is_available(): + major, _ = torch.cuda.get_device_capability() + return major >= 9 # Hopper Compute Capability == 9.0 + return False +``` + +**Rationale:** +- Device names vary across vendors and regions (H100, H800, H20, etc.) +- Compute Capability is a reliable, standardized way to identify GPU architecture +- All Hopper GPUs report `major=9` regardless of their marketing name + +### 2. FA3 Return Value Handling + +Flash Attention 3's `flash_attn_varlen_func` has a different return signature than FA2: + +| Version | Return Value | +|---------|--------------| +| FA2 | `(output, softmax_lse, ...)` - tuple, use `[0]` to get output | +| FA3 | `output` - tensor directly | + +The code correctly handles this difference: + +```python +# FA3 path - direct tensor return +x = flash_attn_interface.flash_attn_varlen_func(...).unflatten(0, (b, lq)) + +# FA2 path - tuple return (handled in else branch) +x = flash_attn.flash_attn_varlen_func(...).unflatten(0, (b, lq)) +``` + +### 3. Automatic Fallback + +The system gracefully falls back to FA2 when FA3 is unavailable: +- If `flash_attn_interface` module is not installed +- If running on non-Hopper GPU +- If user explicitly requests FA2 via `version=2` parameter + +A warning is issued when FA3 is explicitly requested but unavailable. + +## Usage + +### Automatic Selection (Recommended) + +By default, LongLive automatically selects the optimal attention implementation: + +```python +from wan_5b.modules.attention import attention + +# FA3 will be used on Hopper GPUs, FA2 otherwise +output = attention(q, k, v) +``` + +### Explicit Version Selection + +You can force a specific Flash Attention version: + +```python +# Force FA2 (useful for debugging or compatibility) +output = attention(q, k, v, fa_version=2) + +# Request FA3 (falls back to FA2 with warning if unavailable) +output = attention(q, k, v, fa_version=3) +``` + +## Installation + +### Flash Attention 3 (Hopper GPUs) + +```bash +git clone https://github.com/Dao-AILab/flash-attention.git +cd flash-attention/hopper +python setup.py install +``` diff --git a/docs/getting_started.md b/docs/getting_started.md new file mode 100644 index 0000000..b157762 --- /dev/null +++ b/docs/getting_started.md @@ -0,0 +1,430 @@ +# LongLive2.0 Usage + +This document contains the release commands for installation, training, inference, and utilities. The root README keeps the project overview and paper figures. + +## Installation + +Create a Python 3.10 environment and install the required packages: + +```bash +conda create -n longlive2 python=3.10 -y +conda activate longlive2 +pip install torch==2.8.0 torchvision==0.23.0 --index-url https://download.pytorch.org/whl/cu128 +pip install -r requirements.txt +pip install flash-attn --no-build-isolation +``` + +TensorRT is not required for the default training or inference path. If a +TensorRT utility is needed, install it separately after the base requirements: + +```bash +pip install nvidia-pyindex +pip install nvidia-tensorrt +pip install pycuda +``` + +Download the Wan2.2-TI2V-5B components and replace the `/path/to/longlive2.0/...` placeholders in the config files before running training or inference. + +If you set `inference.vae_type` to `mg_lightvae` or `mg_lightvae_v2`, download +the corresponding VAE checkpoints from the Hugging Face repository +`Skywork/Matrix-Game-3.0` and place them under `wan_models/Matrix-Game-3.0/`: + +```text +wan_models/Matrix-Game-3.0/MG-LightVAE.pth +wan_models/Matrix-Game-3.0/MG-LightVAE_v2.pth +``` + +### NVFP4 Environment + +The default installation above is the clean BF16 release setup. NVFP4 training +and inference use local CUDA extensions and are more version-sensitive, so keep +them in a separate environment. + +Known-good NVFP4 baseline inherited from the Sage branch: + +```text +Python: 3.12.12 +PyTorch: 2.10.0+cu128 +TorchVision: 0.25.0+cu128 +CUDA target: 12.8 +FlashAttention: 2.8.3, built from source +``` + +Create or activate the NVFP4 environment: + +```bash +conda create -n longlive2_nvfp4 python=3.12 -y +conda activate longlive2_nvfp4 + +conda install -c nvidia cuda-toolkit=12.8 -y + +pip install -r requirements.txt + +pip install --upgrade --index-url https://download.pytorch.org/whl/cu128 \ + torch==2.10.0 torchvision==0.25.0 +pip install --upgrade torchao==0.16.0 +``` + +If you already have a working `qlive` environment from LongLive_Sage, you can +activate it instead of creating `longlive2_nvfp4`. + +Verify the Torch/CUDA pair: + +```bash +python -c "import torch, torchvision; print(torch.__version__, torch.version.cuda); print(torchvision.__version__)" +``` + +Build the modified local `fouroversix` package: + +```bash +cd fouroversix +pip install ninja packaging psutil "setuptools>=77.0.3" + +# Optional: limit compile targets. +export CUDA_ARCHS=100 # B200 / GB200 / GB300 +# export CUDA_ARCHS=120 # RTX 50/60 series, if needed + +pip install --no-build-isolation -e . +cd .. +``` + +Build FlashAttention from source, rather than relying on a prebuilt wheel: + +```bash +git clone https://github.com/Dao-AILab/flash-attention.git +cd flash-attention +git checkout v2.8.3 +pip install -U pip setuptools wheel ninja packaging +pip install --no-build-isolation -e . +cd .. +``` + +Install TransformerEngine if `model_quant_use_transformer_engine: true` will be +used: + +```bash +python -m pip install --no-build-isolation "transformer-engine[pytorch]" +``` + +Build the fused LongLive FP4 KV-cache dequant extension: + +```bash +cd utils/kernel +python setup.py build_ext --inplace +cd ../.. +``` + +Quick NVFP4 checks: + +```bash +python -c "import flash_attn; print(flash_attn.__version__)" +python -c "import fouroversix; from utils.quant import LongLiveQuantizationConfig, quantize_to_fp4" +python -c "from utils.kernel.kv_dequant import dequantize_kv_cache_fp4" +``` + +The release NVFP4 configs and direct run commands are summarized below. See +`README_NVFP4.md` for lower-level implementation notes. + +## Configs + +The release keeps three main configs: + +```text +configs/train_ar.yaml # AR diffusion training +configs/train_dmd.yaml # DMD distillation +configs/inference.yaml # inference +``` + +TorchAO FP8 PTQ inference has a separate config: + +```text +configs/fp8/inference_fp8.yaml +``` + +The NVFP4 path keeps its configs separate from the default BF16 release path: + +```text +configs/nvfp4/train_ar_nvfp4.yaml # stage 1 AR teacher-forcing training +configs/nvfp4/train_dmd_nvfp4_step4.yaml # stage 2 DMD LoRA distillation, 4-step rollout +configs/nvfp4/inference_nvfp4.yaml # NVFP4 inference with optional KV quantization +``` + +The configs use a shared organization: + +- `model_kwargs`: arguments passed to `WanDiffusionWrapper`. +- `infra`: distributed training/runtime settings. +- `algorithm`: AR or DMD objective settings. +- `training`: optimizer, batch size, checkpoint cadence, and loop settings. +- `data`: training or prompt data paths. +- `inference`: sampling and cache settings. +- `checkpoints`: model and LoRA checkpoint paths. +- `adapter`: optional LoRA settings. Remove this section to disable LoRA. + +## Training + +### AR Diffusion Training + +Edit `configs/train_ar.yaml` to set the dataset path, evaluation prompt path, logging path, and distributed runtime settings. Then run: + +```bash +torchrun --standalone --nnodes=1 --nproc_per_node=8 train.py \ + --config_path configs/train_ar.yaml \ + --logdir logs/test_train_ar \ + --wandb-save-dir wandb \ + --disable-wandb +``` + +Notes: + +- `infra.sequence_parallel_size` controls the SP group size. +- `infra.vae_halo_latents` controls chunk-halo VAE preparation. +- `model_kwargs.local_attn_size` is a model construction setting. +- `inference.sink_size`, `inference.multi_shot_sink`, and `inference.multi_shot_rope_offset` control evaluation-time generation during training. + +### DMD Distillation + +Edit `configs/train_dmd.yaml` to set the dataset path and initialization checkpoints. Then run: + +```bash +torchrun --standalone --nnodes=1 --nproc_per_node=8 train.py \ + --config_path configs/train_dmd.yaml \ + --logdir logs/test_train_dmd \ + --wandb-save-dir wandb \ + --disable-wandb +``` + +Notes: + +- `algorithm.real_guidance_scale` and `algorithm.fake_guidance_scale` are used by score distillation. +- `inference.sampling_steps` controls the distillation rollout sampling steps. +- If `adapter` is present, LoRA distillation is enabled. Otherwise the generator is fully fine-tuned. +- Auto-resume is enabled by default unless `--no-auto-resume` is passed. + +### NVFP4 Training + +Use the `longlive2_nvfp4` environment and build the NVFP4 extensions before +running these commands. Replace the `/path/to/...` placeholders in the configs +first. + +Stage 1 trains the NVFP4 AR teacher-forcing model: + +```bash +torchrun --standalone --nnodes=1 --nproc_per_node=4 train.py \ + --config_path configs/nvfp4/train_ar_nvfp4.yaml \ + --logdir logs/nvfp4_ar \ + --wandb-save-dir wandb \ + --disable-wandb +``` + +Stage 2 runs NVFP4 DMD LoRA distillation from the AR checkpoint: + +```bash +torchrun --standalone --nnodes=1 --nproc_per_node=4 train.py \ + --config_path configs/nvfp4/train_dmd_nvfp4_step4.yaml \ + --logdir logs/nvfp4_dmd_step4 \ + --wandb-save-dir wandb \ + --disable-wandb +``` + +Notes: + +- `--nproc_per_node` controls the per-node GPU count. The NVFP4 examples use 4 +GPUs; set it to 8 or another value for your machine. +- `infra.model_quant` enables NVFP4 generator training for stage 1. +- `infra.generator_quant`, `infra.real_score_quant`, and +`infra.fake_score_quant` choose which DMD networks use NVFP4 in stage 2. + +After stage 1 and stage 2 are complete, you can pre-merge the AR generator and +DMD LoRA weights for inference. The export script reads `generator_ckpt`, +`lora_ckpt`, `adapter`, and `model_quant_*` from the NVFP4 inference config. + +To save a compact FourOverSix materialized NVFP4 generator checkpoint: + +```bash +python scripts/save_merged_nvfp4_generator.py \ + --config_path configs/nvfp4/inference_nvfp4.yaml \ + --output_path /path/to/model_4o6.pt \ + --backend fouroversix \ + --device cuda:0 +``` + +To save merged BF16 weights for TransformerEngine runtime quantization: + +```bash +python scripts/save_merged_nvfp4_generator.py \ + --config_path configs/nvfp4/inference_nvfp4.yaml \ + --output_path /path/to/model_te.pt \ + --backend transformer_engine \ + --device cuda:0 +``` + +The `fouroversix` export is the small packed/materialized NVFP4 checkpoint. The +`transformer_engine` export intentionally saves merged BF16 weights, because a +TransformerEngine module `state_dict` is not a compact packed NVFP4 storage +format; TE quantization is applied again when inference loads the BF16 weights. + +### Merge Generator and LoRA Weights + +For the regular BF16 release path, you can pre-merge the AR generator checkpoint +and DMD LoRA checkpoint into one reusable generator checkpoint. This keeps +quick-start inference simple: inference only loads `checkpoints.generator_ckpt` +and does not need to construct or load LoRA adapters at runtime. + +```bash +python scripts/merge_lora_generator.py \ + --config_path configs/inference.yaml \ + --output_path /path/to/longlive2_merged_generator.pt \ + --device cuda:0 +``` + +After the merge, set `checkpoints.generator_ckpt` in `configs/inference.yaml` to +the merged checkpoint. If you run the full `inference.py` entry point, remove or +unset `checkpoints.lora_ckpt` and the `adapter` section so LoRA is not applied a +second time. + +## Inference + +Edit `configs/inference.yaml` to set: + +- `data.data_path`: prompt folder. +- `checkpoints.generator_ckpt`: merged generator checkpoint. +- `output_folder`: output video directory. +- `num_samples`: number of sampled videos per prompt. + +Run: + +```bash +torchrun --standalone --nnodes=1 --nproc_per_node=8 inference.py \ + --config_path configs/inference.yaml +``` + +Inference notes: + +- `inference.sampling_steps` controls the number of denoising steps. +- `inference.guidance_scale` controls inference CFG. +- `inference.sink_size` controls the standard attention sink size. +- `inference.multi_shot_sink` enables the multi-shot attention sink. +- `inference.multi_shot_rope_offset` controls the multi-shot RoPE offset. + +### FP8 PTQ Inference + +Set `checkpoints.generator_ckpt` in `configs/fp8/inference_fp8.yaml` to the +downloaded merged BF16 `model_bf16.pt`, then run: + +```bash +python inference.py --config_path configs/fp8/inference_fp8.yaml +``` + +`fp8_quant: true` applies TorchAO row-wise dynamic W8A8 quantization after the +generator has been loaded and converted to BF16, and before `torch.compile`. +It cannot be combined with `model_quant: true`, which selects the NVFP4 path. +With the provided 5B model, 300 eligible core Linear layers use FP8 while six +small conditioning/output projections remain BF16 for stability and to avoid +FP8 overhead. + +The validated stack is Python 3.10, PyTorch 2.8.0+cu128, and TorchAO 0.13.0 on +H100 (SM90); compute capability 8.9 or newer is required. The supplied config +uses `torch_compile: auto`: it skips compilation when `inference_iter` +explicitly limits the run to fewer than three samples, and enables it when all +prompts are requested. Its `max-autotune` warm-up can take several minutes +while guard/shape variants are compiled. Use repeated inference and discard all +compile/warm-up samples when measuring steady-state performance; set +`torch_compile: false` for a short eager-mode smoke test. + +The supplied config uses the single 8-latent-frame block validated on H100. +Longer generation introduces additional KV-cache shapes and may trigger more +compilation or eager fallback; validate the intended frame count before +benchmarking or deployment. + +The initial FP8 path targets `inference.py`; `inference_sp.py` rejects the flag +until TorchAO tensor-subclass behavior is validated with Ulysses collectives. + +### NVFP4 Inference + +Edit `configs/nvfp4/inference_nvfp4.yaml` to set: + +- `data.data_path`: prompt folder. +- `checkpoints.generator_ckpt`: AR or base generator checkpoint. +- `checkpoints.lora_ckpt`: optional DMD LoRA checkpoint. +- `output_folder`: output video directory. +- `num_samples`: number of sampled videos per prompt. + +Run: + +```bash +torchrun --standalone --nnodes=1 --nproc_per_node=4 inference.py \ + --config_path configs/nvfp4/inference_nvfp4.yaml +``` + +For single-GPU inference, use `python` directly: + +```bash +python inference.py --config_path configs/nvfp4/inference_nvfp4.yaml +``` + +There are two recommended checkpoint styles for NVFP4 inference: + +FourOverSix compact/materialized NVFP4 checkpoint: + +```yaml +checkpoints: + generator_ckpt: /path/to/model_4o6.pt + +merge_lora: false +model_quant: true +model_quant_use_transformer_engine: false +``` + +TransformerEngine runtime quantization from merged BF16 weights: + +```yaml +checkpoints: + generator_ckpt: /path/to/model_te.pt + +merge_lora: false +model_quant: true +model_quant_use_transformer_engine: true +``` + +Do not set `model_quant_use_transformer_engine: true` when loading a FourOverSix +materialized checkpoint. FourOverSix checkpoints store `quantized_weight_*` +buffers and can only be loaded by the FourOverSix path. TransformerEngine +inference should load merged BF16 weights and quantize them at runtime. + +NVFP4 inference notes: + +- `model_quant` enables generator NVFP4 inference. For regular BF16 +checkpoints, it quantizes/materializes weights during startup; for pre-saved +FourOverSix checkpoints, the checkpoint already contains materialized weights. +- `merge_lora` merges the LoRA checkpoint into the base generator before +quantized materialization. Set it to `false` when `generator_ckpt` already +points to a merged export from `scripts/save_merged_nvfp4_generator.py`. +- `inference.kv_quant` enables FP4 KV-cache storage; the fused dequant extension +from `utils/kernel` must be built first. +- `inference.streaming_vae`, `inference.async_vae`, `inference.vae_type`, and +`inference.vae_device` control streaming or asynchronous VAE decode. +- `torch_compile` can be set to `auto`, `true`, or `false`; the default config +uses `auto` with safe error suppression. + +### Sequence-parallel (SP) inference + +`inference_sp.py` drives **Ulysses sequence-parallel** sampling for WAN (see `configs/inference_sp.yaml` for `sp_size`, `dp_size`, prompts, checkpoints, and the usual `inference.*` knobs). Launch one process per GPU with **`--nproc_per_node` equal to `sp_size × dp_size`** (the shipped example sets `sp_size: 4` and `dp_size: 1`, so four ranks). + +```bash +torchrun --nproc_per_node=4 inference_sp.py --config_path configs/inference_sp.yaml +``` + +## Utilities + +Inspect SP VAE halo windows: + +```bash +python scripts/compute_sp_vae_chunk_halo.py --config configs/train_ar.yaml +``` + +Decode saved VAE latents: + +```bash +python scripts/decode_vae_latents.py --help +python scripts/decode_lightvae_latents.py --help +``` diff --git a/example/long_example.txt b/example/long_example.txt new file mode 100644 index 0000000..0865419 --- /dev/null +++ b/example/long_example.txt @@ -0,0 +1,8 @@ +A close-up shot of a ceramic teacup slowly pouring water into a glass mug. The water flows smoothly from the spout of the teacup into the mug, creating gentle ripples as it fills up. Both cups have detailed textures, with the teacup having a matte finish and the glass mug showcasing clear transparency. The background is a blurred kitchen countertop, adding context without distracting from the central action. The pouring motion is fluid and natural, emphasizing the interaction between the two cups. +A dynamic and chaotic scene in a dense forest during a heavy rainstorm, capturing a real girl frantically running through the foliage. Her wild hair flows behind her as she sprints, her arms flailing and her face contorted in fear and desperation. Behind her, various animals—rabbits, deer, and birds—are also running, creating a frenzied atmosphere. The girl's clothes are soaked, clinging to her body, and she is screaming and shouting as she tries to escape. The background is a blur of greenery and rain-drenched trees, with occasional glimpses of the darkening sky. A wide-angle shot from a low angle, emphasizing the urgency and chaos of the moment. +A dynamic over-the-shoulder perspective of a chef meticulously plating a dish in a bustling kitchen. The chef, a middle-aged man with a neatly trimmed beard and focused expression, deftly arranges ingredients on a pristine white plate. His hands move with precision, each gesture deliberate and practiced. The background shows a crowded kitchen with steaming pots, whirring blenders, and the clatter of utensils. Bright lights highlight the scene, casting shadows across the busy workspace. The camera angle captures the chef's detailed work from behind, emphasizing his skill and dedication. +A dramatic and dynamic scene in the style of a disaster movie, depicting a powerful tsunami rushing through a narrow alley in Bulgaria. The water is turbulent and chaotic, with waves crashing violently against the walls and buildings on either side. The alley is lined with old, weathered houses, their facades partially submerged and splintered. The camera angle is low, capturing the full force of the tsunami as it surges forward, creating a sense of urgency and danger. People can be seen running frantically, adding to the chaos. The background features a distant horizon, hinting at the larger scale of the tsunami. A dynamic, sweeping shot from a low-angle perspective, emphasizing the movement and intensity of the event. +A playful raccoon is seen playing an electronic guitar, strumming the strings with its front paws. The raccoon has distinctive black facial markings and a bushy tail. It sits comfortably on a small stool, its body slightly tilted as it focuses intently on the instrument. The setting is a cozy, dimly lit room with vintage posters on the walls, adding a retro vibe. The raccoon's expressive eyes convey a sense of joy and concentration. Medium close-up shot, focusing on the raccoon's face and hands interacting with the guitar. +A single white sheep bending down to drink water from a calm river. The sheep has fluffy wool, long curved horns, and soft brown eyes. It is positioned near the riverbank, with its front legs partially submerged in the clear water. The river flows gently, reflecting the surrounding greenery and blue sky. The background shows lush grass and trees along the riverbank, creating a serene pastoral landscape. The sheep's body is slightly tilted as it bends down to drink, emphasizing a natural and tranquil motion. Medium close-up shot focusing on the sheep and the river. +3D animation of a small, round, fluffy creature with big, expressive eyes explores a vibrant, enchanted forest. The creature, a whimsical blend of a rabbit and a squirrel, has soft blue fur and a bushy, striped tail. It hops along a sparkling stream, its eyes wide with wonder. The forest is alive with magical elements: flowers that glow and change colors, trees with leaves in shades of purple and silver, and small floating lights that resemble fireflies. The creature stops to interact playfully with a group of tiny, fairy-like beings dancing around a mushroom ring. The creature looks up in awe at a large, glowing tree that seems to be the heart of the forest. +Animated scene features a close-up of a short fluffy monster kneeling beside a melting red candle. The art style is 3D and realistic, with a focus on lighting and texture. The mood of the painting is one of wonder and curiosity, as the monster gazes at the flame with wide eyes and open mouth. Its pose and expression convey a sense of innocence and playfulness, as if it is exploring the world around it for the first time. The use of warm colors and dramatic lighting further enhances the cozy atmosphere of the image. \ No newline at end of file diff --git a/fouroversix/.github/workflows/_build.yml b/fouroversix/.github/workflows/_build.yml new file mode 100644 index 0000000..3ae1354 --- /dev/null +++ b/fouroversix/.github/workflows/_build.yml @@ -0,0 +1,215 @@ +name: ~Build wheel template + +on: + workflow_call: + inputs: + runs-on: + description: "The runner to use for the build" + required: true + type: string + python-version: + description: "The Python version to use for the build" + required: true + type: string + cuda-version: + description: "The CUDA version to use for the build" + required: true + type: string + torch-version: + description: "The PyTorch version to use for the build" + required: true + type: string + cxx11_abi: + description: "The C++11 ABI to use for the build" + required: true + type: string + upload-to-release: + description: "Upload wheel to this release" + required: false + type: boolean + default: false + release-version: + description: "Upload wheel to this release" + required: false + type: string + +defaults: + run: + shell: bash -x -e -u -o pipefail {0} + +jobs: + build-wheel: + runs-on: ${{ inputs.runs-on }} + name: Build wheel (${{ inputs.release-version }}-${{ inputs.python-version }}-${{ inputs.cuda-version }}-${{ inputs.torch-version }}-${{ inputs.cxx11_abi }}) + steps: + - name: Checkout + uses: actions/checkout@v5 + with: + ref: ${{ inputs.release-version }} + submodules: recursive + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: ${{ inputs.python-version }} + + - name: Set CUDA and PyTorch versions + run: | + echo "MATRIX_CUDA_VERSION=$(echo ${{ inputs.cuda-version }} | awk -F \. {'print $1 $2'})" >> $GITHUB_ENV + echo "MATRIX_TORCH_VERSION=$(echo ${{ inputs.torch-version }} | awk -F \. {'print $1 "." $2'})" >> $GITHUB_ENV + echo "WHEEL_CUDA_VERSION=$(echo ${{ inputs.cuda-version }} | awk -F \. {'print $1'})" >> $GITHUB_ENV + echo "MATRIX_PYTHON_VERSION=$(echo ${{ inputs.python-version }} | awk -F \. {'print $1 $2'})" >> $GITHUB_ENV + + - name: Free up disk space + if: ${{ runner.os == 'Linux' }} + # https://github.com/easimon/maximize-build-space/blob/master/action.yml + # https://github.com/easimon/maximize-build-space/tree/test-report + run: | + sudo rm -rf /usr/share/dotnet + sudo rm -rf /usr/local/lib/android + sudo rm -rf /opt/ghc + sudo rm -rf /opt/hostedtoolcache/CodeQL + sudo docker image prune --all --force + + - name: Set up swap space + if: runner.os == 'Linux' + uses: pierotofy/set-swap-space@v1.0 + with: + swap-size-gb: 10 + + - name: Install CUDA ${{ inputs.cuda-version }} + if: ${{ inputs.cuda-version != 'cpu' }} + uses: Jimver/cuda-toolkit@v0.2.29 + id: cuda-toolkit + with: + cuda: ${{ inputs.cuda-version }} + linux-local-args: '["--toolkit"]' + # default method is "local", and we're hitting some error with caching for CUDA 11.8 and 12.1 + # method: ${{ (inputs.cuda-version == '11.8.0' || inputs.cuda-version == '12.1.0') && 'network' || 'local' }} + method: "network" + sub-packages: '["nvcc"]' + non-cuda-sub-packages: '["libcublas-dev", "libcusolver-dev", "libcusparse-dev"]' + + - name: Install PyTorch ${{ inputs.torch-version }}+cu${{ inputs.cuda-version }} + run: | + pip install --upgrade pip + # With python 3.13 and torch 2.5.1, unless we update typing-extensions, we get error + # AttributeError: attribute '__default__' of 'typing.ParamSpec' objects is not writable + pip install typing-extensions==4.12.2 + # detect if we're on ARM + if [ "$(uname -m)" = "aarch64" ] || [ "$(uname -m)" = "arm64" ]; then + PLAT=linux_aarch64 + else + PLAT=manylinux_2_27_x86_64.manylinux_2_28_x86_64 + fi + echo "PLAT=$PLAT" >> $GITHUB_ENV + if [[ ${{ inputs.torch-version }} == *"dev"* ]]; then + # pip install --no-cache-dir --pre torch==${{ inputs.torch-version }} --index-url https://download.pytorch.org/whl/nightly/cu${MATRIX_CUDA_VERSION} + # Can't use --no-deps because we need cudnn etc. + # Hard-coding this version of pytorch-triton for torch 2.9.0.dev20250904 + pip install jinja2 + TRITON_URL=https://download.pytorch.org/whl/nightly/pytorch_triton-3.4.0%2Bgitf7888497-cp${MATRIX_PYTHON_VERSION}-cp${MATRIX_PYTHON_VERSION}-${PLAT}.whl + TORCH_URL=https://download.pytorch.org/whl/nightly/cu${MATRIX_CUDA_VERSION}/torch-${{ inputs.torch-version }}%2Bcu${MATRIX_CUDA_VERSION}-cp${MATRIX_PYTHON_VERSION}-cp${MATRIX_PYTHON_VERSION}-manylinux_2_28_$(uname -m).whl + pip install --no-cache-dir --pre "${TRITON_URL}" + pip install --no-cache-dir --pre "${TORCH_URL}" + else + pip install --no-cache-dir torch==${{ inputs.torch-version }} --index-url https://download.pytorch.org/whl/cu${MATRIX_CUDA_VERSION} + fi + nvcc --version + python --version + python -c "import torch; print('PyTorch:', torch.__version__)" + python -c "import torch; print('CUDA:', torch.version.cuda)" + python -c "from torch.utils import cpp_extension; print (cpp_extension.CUDA_HOME)" + + - name: Restore build cache + uses: actions/cache/restore@v4 + with: + path: build.tar + key: build-${{ inputs.release-version }}-${{ inputs.python-version }}-${{ inputs.cuda-version }}-${{ inputs.torch-version }}-${{ inputs.cxx11_abi }}-${{ github.run_number }}-${{ github.run_attempt }} + restore-keys: | + build-${{ inputs.release-version }}-${{ inputs.python-version }}-${{ inputs.cuda-version }}-${{ inputs.torch-version }}-${{ inputs.cxx11_abi }}- + + - name: Unpack build cache + run: | + echo ::group::Adjust timestamps + sudo find / -exec touch -t 197001010000 {} + || true + echo ::endgroup:: + + if [ -f build.tar ]; then + find . -mindepth 1 -maxdepth 1 ! -name 'build.tar' -exec rm -rf {} + + tar -xpvf build.tar -C . + else + echo "No build.tar found, skipping" + fi + + ls -al ./ + ls -al build/ || true + ls -al csrc/ || true + + - name: Build wheel + id: build_wheel + run: | + pip install -U ninja packaging psutil setuptools wheel + export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH + export LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/usr/local/cuda/lib64:$LD_LIBRARY_PATH + # Limit MAX_JOBS otherwise the github runner goes OOM + # nvcc 11.8 can compile with 2 jobs, but nvcc 12.3 goes OOM + + export MAX_JOBS=$([ "$MATRIX_CUDA_VERSION" == "129" ] && echo 1 || echo 2) + export NVCC_THREADS=2 + export FORCE_BUILD=1 + export FORCE_CXX11_ABI=${{ inputs.cxx11_abi == 'TRUE' && '1' || '0' }} + + # 5h timeout since GH allows max 6h and we want some buffer + EXIT_CODE=0 + timeout 5h python setup.py bdist_wheel --dist-dir=dist || EXIT_CODE=$? + + if [ $EXIT_CODE -eq 0 ]; then + tmpname=cu${WHEEL_CUDA_VERSION}torch${MATRIX_TORCH_VERSION}cxx11abi${{ inputs.cxx11_abi }} + wheel_name=$(ls dist/*whl | xargs -n 1 basename | sed "s/-/+$tmpname-/2") + ls dist/*whl |xargs -I {} mv {} dist/${wheel_name} + echo "wheel_name=${wheel_name}" >> $GITHUB_ENV + fi + + # Store exit code in GitHub env for later steps + echo "build_exit_code=$EXIT_CODE" | tee -a "$GITHUB_OUTPUT" + + # Do not fail the job if timeout killed the build + exit $EXIT_CODE + + - name: Log build logs after timeout + if: always() && steps.build_wheel.outputs.build_exit_code == 124 + run: | + ls -al ./ + tar -cvf build.tar . --atime-preserve=replace + + - name: Save build cache timeout + if: always() && steps.build_wheel.outputs.build_exit_code == 124 + uses: actions/cache/save@v4 + with: + key: build-${{ inputs.release-version }}-${{ inputs.python-version }}-${{ inputs.cuda-version }}-${{ inputs.torch-version }}-${{ inputs.cxx11_abi }}-${{ github.run_number }}-${{ github.run_attempt }} + path: build.tar + + - name: Log Built Wheels + run: | + ls dist + + - name: Get Release with tag + id: get_current_release + uses: joutvhu/get-release@v1 + with: + tag_name: ${{ inputs.release-version }} + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + + - name: Upload Release Asset + id: upload_release_asset + if: inputs.upload-to-release + uses: actions/upload-release-asset@v1 + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + with: + upload_url: ${{ steps.get_current_release.outputs.upload_url }} + asset_path: ./dist/${{env.wheel_name}} + asset_name: ${{env.wheel_name}} + asset_content_type: application/* diff --git a/fouroversix/.github/workflows/build.yml b/fouroversix/.github/workflows/build.yml new file mode 100644 index 0000000..25ea5e8 --- /dev/null +++ b/fouroversix/.github/workflows/build.yml @@ -0,0 +1,47 @@ +name: Build wheels + +on: + workflow_dispatch: + inputs: + runs-on: + description: "The runner to use for the build" + required: true + type: string + default: ubuntu-22.04 + python-version: + description: "The Python version to use for the build" + required: true + type: string + cuda-version: + description: "The CUDA version to use for the build" + required: true + type: string + torch-version: + description: "The PyTorch version to use for the build" + required: true + type: string + cxx11_abi: + description: "Enable torch flag C++11 ABI (TRUE/FALSE)" + required: true + type: string + upload-to-release: + description: "Upload wheel to this release" + required: false + type: boolean + default: false + release-version: + description: "Upload wheel to this release" + required: false + type: string + +jobs: + build-wheels: + uses: ./.github/workflows/_build.yml + with: + runs-on: ${{ inputs.runs-on }} + python-version: ${{ inputs.python-version }} + cuda-version: ${{ inputs.cuda-version }} + torch-version: ${{ inputs.torch-version }} + cxx11_abi: ${{ inputs.cxx11_abi }} + upload-to-release: ${{ inputs.upload-to-release }} + release-version: ${{ inputs.release-version }} diff --git a/fouroversix/.github/workflows/pre-commit.yml b/fouroversix/.github/workflows/pre-commit.yml new file mode 100644 index 0000000..d8b5fb6 --- /dev/null +++ b/fouroversix/.github/workflows/pre-commit.yml @@ -0,0 +1,16 @@ +name: Lint + +on: + push: + branches: + - main + +jobs: + lint: + runs-on: ubuntu-latest + strategy: + matrix: + python: ["3.13"] + steps: + - uses: actions/checkout@v4 + - uses: astral-sh/ruff-action@v3 diff --git a/fouroversix/.github/workflows/publish.yml b/fouroversix/.github/workflows/publish.yml new file mode 100644 index 0000000..eb5b172 --- /dev/null +++ b/fouroversix/.github/workflows/publish.yml @@ -0,0 +1,101 @@ +# This workflow will: +# - Create a new Github release +# - Build wheels for supported architectures +# - Deploy the wheels to the Github release +# - Release the static code to PyPi +# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries + +name: Build wheels and deploy + +on: + push: + tags: + - "v*" + +jobs: + setup_release: + name: Create Release + runs-on: ubuntu-latest + outputs: + release-version: ${{ steps.extract_branch.outputs.branch }} + steps: + - name: Get the tag version + id: extract_branch + run: echo ::set-output name=branch::${GITHUB_REF#refs/tags/} + shell: bash + - name: Create Release + id: create_release + uses: actions/create-release@v1 + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + with: + tag_name: ${{ steps.extract_branch.outputs.branch }} + release_name: ${{ steps.extract_branch.outputs.branch }} + + build_wheels: + name: Build Wheel + needs: setup_release + strategy: + fail-fast: false + matrix: + # Using ubuntu-22.04 instead of 24.04 for more compatibility (glibc). Ideally we'd use the + # manylinux docker image, but I haven't figured out how to install CUDA on manylinux. + os: [ubuntu-22.04, ubuntu-22.04-arm] + python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] + torch-version: ["2.8.0", "2.9.1", "2.10.0"] + cuda-version: ["12.9.1", "13.0.2"] + # We need separate wheels that either uses C++11 ABI (-D_GLIBCXX_USE_CXX11_ABI) or not. + # Pytorch wheels currently don't use it, but nvcr images have Pytorch compiled with C++11 ABI. + # Without this we get import error (undefined symbol: _ZN3c105ErrorC2ENS_14SourceLocationESs) + # when building without C++11 ABI and using it on nvcr images. + cxx11_abi: ["FALSE", "TRUE"] + exclude: + - os: ubuntu-22.04-arm + torch-version: "2.8.0" + - os: ubuntu-22.04-arm + torch-version: "2.7.1" + # PyTorch 2.8 only supports up to CUDA 12.9 + - torch-version: "2.8.0" + cuda-version: "13.0.2" + - torch-version: "2.7.1" + cuda-version: "13.0.2" + # Python 3.14 only has pre-built wheels for PyTorch 2.9 and newer + - python-version: "3.14" + torch-version: "2.7.1" + - python-version: "3.14" + torch-version: "2.8.0" + uses: ./.github/workflows/_build.yml + with: + runs-on: ${{ matrix.os }} + python-version: ${{ matrix.python-version }} + cuda-version: ${{ matrix.cuda-version }} + torch-version: ${{ matrix.torch-version }} + cxx11_abi: ${{ matrix.cxx11_abi }} + release-version: ${{ needs.setup_release.outputs.release-version }} + upload-to-release: true + + publish_package: + name: Publish package + needs: [build_wheels] + runs-on: ubuntu-latest + permissions: + id-token: write + steps: + - uses: actions/checkout@v5 + with: + submodules: recursive + - uses: actions/setup-python@v6 + with: + python-version: "3.10" + - name: Install dependencies + run: | + pip install -U ninja packaging setuptools wheel twine + # We don't want to download anything CUDA-related here + pip install torch --index-url https://download.pytorch.org/whl/cpu + - name: Build core package + env: + SKIP_CUDA_BUILD: "1" + run: | + python setup.py sdist --dist-dir=dist + - name: Deploy + uses: pypa/gh-action-pypi-publish@release/v1 diff --git a/fouroversix/.gitignore b/fouroversix/.gitignore new file mode 100644 index 0000000..65af597 --- /dev/null +++ b/fouroversix/.gitignore @@ -0,0 +1,16 @@ +__pycache__ +*.py[cod] +.DS_Store +.vscode +env/ +trace_*.json +*.so +exp/ +*.ipynb +.cache/ +ptq_logs/ +dist/ +*.egg-info/ +site/ +results.db +/ptq/ \ No newline at end of file diff --git a/fouroversix/.gitmodules b/fouroversix/.gitmodules new file mode 100644 index 0000000..feb3c5f --- /dev/null +++ b/fouroversix/.gitmodules @@ -0,0 +1,21 @@ +[submodule "third_party/cutlass"] + path = third_party/cutlass + url = https://github.com/NVIDIA/cutlass.git +[submodule "third_party/fp-quant"] + path = third_party/fp-quant + url = https://github.com/jackcook/fp-quant-fouroversix.git +[submodule "third_party/llm-awq"] + path = third_party/llm-awq + url = https://github.com/jackcook/llm-awq-fouroversix.git +[submodule "third_party/qutlass"] + path = third_party/qutlass + url = https://github.com/IST-DASLab/qutlass.git +[submodule "third_party/fast-hadamard-transform"] + path = third_party/fast-hadamard-transform + url = https://github.com/Dao-AILab/fast-hadamard-transform.git +[submodule "third_party/spinquant"] + path = third_party/spinquant + url = https://github.com/jackcook/spinquant-fouroversix.git +[submodule "third_party/flame"] + path = third_party/flame + url = https://github.com/jackcook/flame-fouroversix.git diff --git a/fouroversix/LICENSE.md b/fouroversix/LICENSE.md new file mode 100644 index 0000000..827e2c0 --- /dev/null +++ b/fouroversix/LICENSE.md @@ -0,0 +1,21 @@ +# MIT License + +Copyright (c) 2025 Jack Cook + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/fouroversix/MANIFEST.in b/fouroversix/MANIFEST.in new file mode 100644 index 0000000..2ecd064 --- /dev/null +++ b/fouroversix/MANIFEST.in @@ -0,0 +1,12 @@ +recursive-include src *.cu +recursive-include src *.h +recursive-include src *.cpp +recursive-include src *.hpp +recursive-include src *.py + +recursive-include third_party/cutlass *.cu +recursive-include third_party/cutlass *.h +recursive-include third_party/cutlass *.cuh +recursive-include third_party/cutlass *.cpp +recursive-include third_party/cutlass *.hpp +recursive-include third_party/cutlass *.inl \ No newline at end of file diff --git a/fouroversix/README.md b/fouroversix/README.md new file mode 100644 index 0000000..3ff117a --- /dev/null +++ b/fouroversix/README.md @@ -0,0 +1,167 @@ +# Four Over Six (4/6) + +[![arXiv](https://img.shields.io/badge/arXiv-2512.02010-b31b1b.svg)](https://arxiv.org/abs/2512.02010) + +_Improving the accuracy of NVFP4 quantization with Adaptive Block Scaling._ + +![](/assets/four-over-six.png) + +This repository contains kernels for efficient NVFP4 quantization and matrix multiplication, and fast post-training quantization with our method, 4/6. +If you have any questions, please get in touch or submit an issue. + +## Setup + +**Requirements:** + +- Python version 3.10 or newer +- CUDA toolkit 12.8 or newer +- PyTorch version 2.8 or newer + +**Install dependencies:** + +```bash +pip install ninja packaging psutil "setuptools>=77.0.3" +``` + +**Install fouroversix:** + +```bash +pip install fouroversix --no-build-isolation +``` + +Alternatively, you can compile from source: + +```bash +pip install --no-build-isolation -e . +``` + +To speed up build times, set `CUDA_ARCHS=100` to only compile kernels for B-series GPUs (i.e. B200, GB200, GB300), or `CUDA_ARCHS=120` for RTX 50 and 60 Series GPUs (i.e. RTX 5090, RTX 6000). + +Also, if you don't have a Blackwell GPU, you may use our reference implementation, which is slow but helpful for testing, by setting `SKIP_CUDA_BUILD=1` before running `pip install`. + +### PTQ Experiments + +To run PTQ experiments, make sure to install our test dependencies using either: + +```bash +pip install "fouroversix[evals]" --no-build-isolation + +# Or, if installing from source: +pip install --no-build-isolation -e ".[evals]" +``` + +Also, make sure all submodules are pulled and up to date: + +```bash +git submodule update --init +``` + +Then, install dependencies for each PTQ method as needed, following the instructions [here](/docs/ptq.md). + +## API + +### Quantize a Model to NVFP4 + +```python +from fouroversix import ModelQuantizationConfig, quantize_model +from transformers import AutoModelForCausalLM + +# NVFP4 using 4/6 with MSE block selection +model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B") +quantize_model(model) + +# Standard NVFP4 round-to-nearest quantization +model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B") +config = ModelQuantizationConfig(scale_rule="static_6") +quantize_model(model, config) +``` + +### Quantize a Tensor to NVFP4 + +Check the `quantize_to_fp4` [arguments](https://github.com/mit-han-lab/fouroversix/blob/f1b78701c753ea49c091ac39d85c5753b703f5ca/src/fouroversix/frontend.py#L72) for more details about how you can enable certain features during quantization, such as stochastic rounding or 2D block quantization. + +```python +import torch +from fouroversix import QuantizationConfig, quantize_to_fp4 + +x = torch.randn(1024, 1024, dtype=torch.bfloat16, device="cuda") +x_quantized = quantize_to_fp4(x) + +# Standard NVFP4 round-to-nearest quantization +config = QuantizationConfig(scale_rule="static_6") +x_quantized = quantize_to_fp4(x, config) +``` + +### Multiply Two NVFP4 Tensors + +```python +from fouroversix import fp4_matmul + +# a and b can be either high-precision BF16 tensors, in which case they will be +# quantized, or low-precision QuantizedTensors if you've already quantized them +# yourself. +out = fp4_matmul(a, b) +``` + +## PTQ Evaluation with LM Evaluation Harness + +```bash +# Round-to-nearest quantization with 4/6: +python -m scripts.ptq --model-name meta-llama/Llama-3.2-1B --ptq-method rtn --task wikitext + +# Standard NVFP4 round-to-nearest (RTN) quantization: +python -m scripts.ptq --model-name meta-llama/Llama-3.2-1B --ptq-method rtn --task wikitext --a-scale-rule static_6 --w-scale-rule static_6 + +# AWQ with 4/6: +python -m scripts.ptq --model-name meta-llama/Llama-3.2-1B --ptq-method awq --task wikitext + +# High-precision baseline, no NVFP4 quantization: +python -m scripts.ptq --model-name meta-llama/Llama-3.2-1B --ptq-method high_precision --task wikitext +``` + +If you would prefer not to worry about setting up your local environment, or about acquiring a Blackwell GPU to run your experiments faster, you may run PTQ experiments on [Modal](https://modal.com/) by adding the `--modal` flag, and optionally the `--detach` flag which will enable you to CTRL+C. +The first time you launch experiments on Modal, it may take several minutes to build everything, but following commands will reuse the cached images. + +## Notes + +This repository contains three implementations of NVFP4 quantization, each of which has various limitations: + +- [CUDA](/src/fouroversix/csrc): Supports most but not all operations needed for efficient NVFP4 training. More operations will be added soon. Requires a Blackwell GPU. +- [Triton](/src/fouroversix/quantize/triton_kernel.py): Supports all operations needed for efficient NVFP4 training, including stochastic rounding, the random Hadamard transform, transposed inputs, and 2D block scaling. Requires a Blackwell GPU. +- [PyTorch](/src/fouroversix/quantize/reference.py): A reference implementation written in PyTorch that can run on any GPU. May have some educational value. Should not be used in real-world use cases. + +When used with 4/6, these implementations have subtle numerical differences which can cause results to differ slightly, but not in a way that should cause uniformly worse performance for any of them. +For more details, see [here](https://github.com/mit-han-lab/fouroversix/blob/6bb13a8fc3b690154d11a1d6477bb6c2d09799e8/tests/test_correctness.py#L124-L132). + +Our `quantize_to_fp4` function will automatically select one of these backends based on your GPU and the quantization parameters you select. +If you would like to force selection of a specific backend, you may specify it by setting `backend=QuantizeBackend.cuda` in the quantization config passed to `quantize_to_fp4`, or `quantize_backend=QuantizeBackend.cuda` in the layer and model configs passed to `quantize_model`. + +## Contributing + +We welcome contributions to our repository, but get in touch before making any substantial changes. +Also, please make sure any code changes are compliant with our linter: + +```bash +ruff check +``` + +## Citation + +Please use the following BibTeX entry to cite this work: + +```bibtex +@misc{cook2025sixaccuratenvfp4quantization, + title={Four Over Six: More Accurate NVFP4 Quantization with Adaptive Block Scaling}, + author={Jack Cook and Junxian Guo and Guangxuan Xiao and Yujun Lin and Song Han}, + year={2025}, + eprint={2512.02010}, + archivePrefix={arXiv}, + primaryClass={cs.CL}, + url={https://arxiv.org/abs/2512.02010}, +} +``` + +## License + +This repository is available under the MIT license. +See the [LICENSE.md](/LICENSE.md) file for details. diff --git a/fouroversix/docs/index.md b/fouroversix/docs/index.md new file mode 100644 index 0000000..4c3c9c2 --- /dev/null +++ b/fouroversix/docs/index.md @@ -0,0 +1,3 @@ +# Four Over Six + +See [Quantization](/quantization) for details on quantizing tensors to FP4, and [Matrix Multiplication](/matmul) for performing matrix multiplication with FP4 tensors. diff --git a/fouroversix/docs/matmul.md b/fouroversix/docs/matmul.md new file mode 100644 index 0000000..dd99950 --- /dev/null +++ b/fouroversix/docs/matmul.md @@ -0,0 +1,3 @@ +# Matrix Multiplication + +::: fouroversix.fp4_matmul diff --git a/fouroversix/docs/ptq.md b/fouroversix/docs/ptq.md new file mode 100644 index 0000000..893affa --- /dev/null +++ b/fouroversix/docs/ptq.md @@ -0,0 +1,69 @@ +# Running PTQ Experiments + +## Install Dependencies + +Before doing anything, make sure you've installed fouroversix with our test dependencies: + +```bash +pip install -e .[evals] --no-build-isolation +``` + +Also, make sure you've cloned all of our submodules: + +```bash +git submodule update --init +``` + +Then, depending on which PTQ method you would like to test, you may need to run some additional commands. + +### AWQ + +```bash +pip install --no-deps third_party/llm-awq +``` + +### GPTQ + +1. Install Fast Hadamard Transform + +```bash +pip install --no-build-isolation third_party/fast-hadamard-transform +``` + +2. Install QuTLASS + +```bash +pip install --no-build-isolation third_party/qutlass +``` + +3. Install FP-Quant + +```bash +pip install third_party/fp-quant/inference_lib +``` + +### High Precision + +No installation necessary! + +### Round-to-Nearest (RTN) + +No installation necessary! + +### SmoothQuant + +No installation necessary! + +### SpinQuant + +1. Install Fast Hadamard Transform + +```bash +pip install --no-build-isolation third_party/fast-hadamard-transform +``` + +2. Downgrade Transformers if your installation is up-to-date + +```bash +pip install "transformers<5.0" +``` \ No newline at end of file diff --git a/fouroversix/docs/quantization.md b/fouroversix/docs/quantization.md new file mode 100644 index 0000000..7d2cf72 --- /dev/null +++ b/fouroversix/docs/quantization.md @@ -0,0 +1,3 @@ +# Quantization + +::: fouroversix.quantize_to_fp4 diff --git a/fouroversix/mkdocs.yml b/fouroversix/mkdocs.yml new file mode 100644 index 0000000..4f5bd1e --- /dev/null +++ b/fouroversix/mkdocs.yml @@ -0,0 +1,8 @@ +site_name: Four Over Six Documentation + +theme: + name: material + +plugins: + - search + - mkdocstrings diff --git a/fouroversix/pyproject.toml b/fouroversix/pyproject.toml new file mode 100644 index 0000000..3677ef9 --- /dev/null +++ b/fouroversix/pyproject.toml @@ -0,0 +1,110 @@ +[build-system] +requires = [ + "ninja", + "packaging", + "psutil", + "setuptools>=77.0.3", +] +build-backend = "setuptools.build_meta" + +[project] +name = "fouroversix" +dynamic = ["version"] +authors = [ + { name="Jack Cook", email="cookj@mit.edu" }, + { name="Junxian Guo", email="junxian@mit.edu" }, +] +description = "More Accurate FP4 Quantization with Adaptive Block Scaling" +readme = "README.md" +requires-python = ">=3.10" +license = "MIT" +license-files = ["LICENSE.md"] +dependencies = [ + "torch>=2.7.0", +] + +[project.optional-dependencies] +docs = [ + "mkdocs~=1.6.1", + "mkdocs-material~=9.7.4", + "mkdocstrings[python]~=1.0.3", +] +evals = [ + "inspect-ai~=0.3.186", + "inspect-evals~=0.3.106", + "lm-eval[hf]~=0.4.11", + "modal~=1.3.5", + "openai~=2.24.0", + "pytest~=8.1.1", + "ruff~=0.15.4", + "SQLAlchemy~=2.0.48", + "transformers @ git+https://github.com/huggingface/transformers.git@22c35ca", +] + +[tool.pytest.ini_options] +testpaths = ["tests"] + +[tool.ruff] +line-length = 88 +indent-width = 4 +exclude = ["build", "dist", "env", ".venv", "third_party", "scripts/ptq/tasks"] + +[tool.ruff.format] +quote-style = "double" +indent-style = "space" + +[tool.ruff.lint] +ignore = [ + # Various dumb rules related to docstrings + "D100", + "D103", # TODO: Remove this one + "D104", + "D107", + "D202", + "D203", + "D205", + "D212", + + "A001", "A002", # Allow variables to be named `input` + "FIX002", # Allow TODO comments + "FIX004", # Allow hacks + "N812", # Allow importing torch.nn.functional as F + "PLR0402", # Allow importing torch.nn as nn + "PLC0415", # Allow imports that aren't at the top of the file + "PLR0913", # Allow functions with more than 5 arguments + "RET506", # Allow elif after raise + "S603", # Allow untrusted inputs to be passed to subprocesses + "S607", # Allow commands to be run by subprocess without providing a full path + "TD003", # Don't require issue links on TODO comments +] +select = [ + "ALL", +] + +[tool.ruff.lint.extend-per-file-ignores] +"setup.py" = [ + "T201", # Allow print statements in setup.py +] +"scripts/**/*.py" = [ + "T201", # Allow print statements in CLI tools + "TID252", # Allow relative imports from parent modules +] +"src/fouroversix/**/ops.py" = [ + "I001", # Allow unsorted imports, because torch needs to be imported before fouroversix._C +] +"src/fouroversix/quantize/pytorch/reference.py" = [ + "S101", # Allow assert statements in reference implementation +] +"src/fouroversix/quantize/triton/kernel.py" = [ + "ANN001", # Allow missing annotations in Triton kernels + "N803", # Allow uppercase arguments to Triton kernels + "N806", # Allow uppercase variables in Triton kernels + "PLR1714", "SIM109", # Allow multiple equality comparisons in Triton kernels +] +"tests/**/test_*.py" = [ + "S101", # Allow assert statements in tests + "T201", # Allow print statements in tests +] + +[tool.ruff.lint.isort] +known-third-party = ["fouroversix"] \ No newline at end of file diff --git a/fouroversix/scripts/__init__.py b/fouroversix/scripts/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/fouroversix/scripts/create_test_case.py b/fouroversix/scripts/create_test_case.py new file mode 100644 index 0000000..2b0121e --- /dev/null +++ b/fouroversix/scripts/create_test_case.py @@ -0,0 +1,87 @@ +from .resources import Dependency, app, get_image + +img = get_image(dependencies=[Dependency.transformer_engine, Dependency.fouroversix]) + +with img.imports(): + import torch + from fouroversix import AdaptiveBlockScalingRule, QuantizeBackend, quantize_to_fp4 + from fouroversix.quantize import from_blocked + + +@app.function(image=img, gpu="B200") +def create_test_case( + backend_a: str = "cuda", + backend_b: str = "transformer_engine", + scale_rule: str = "mse", +) -> None: + M, N = 1024, 1024 # noqa: N806 + + torch.set_printoptions(precision=10) + + backend_a = QuantizeBackend(backend_a) + backend_b = QuantizeBackend(backend_b) + scale_rule = AdaptiveBlockScalingRule(scale_rule) + + for random_seed in range(10): + torch.manual_seed(random_seed) + + x = torch.randn(M, N, dtype=torch.bfloat16, device="cuda") + out_a = quantize_to_fp4( + x, + backend=backend_a, + scale_rule=scale_rule, + ) + out_b = quantize_to_fp4( + x, + backend=backend_b, + scale_rule=scale_rule, + ) + x_sf_a = from_blocked(out_a.scale_factors.bfloat16(), (M, N // 16)) + x_sf_b = from_blocked(out_b.scale_factors.bfloat16(), (M, N // 16)) + + print(f"x absmax: {x.abs().max()}") + + if not torch.allclose(out_a.amax, out_b.amax): + print("Backends A and B have different amax values!") + print(f"{backend_a}: {out_a.amax}") + print(f"{backend_b}: {out_b.amax}") + return + + if not torch.allclose(x_sf_a.bfloat16(), x_sf_b.bfloat16()): + mismatch_prop = (x_sf_a != x_sf_b).sum() / x_sf_a.numel() + print( + "Backends A and B have different scale factors! " + f"{mismatch_prop:.2%} mismatch", + ) + + [i, *_], [j, *_] = torch.where(x_sf_a != x_sf_b) + print(backend_a) + print("sf", x_sf_a[i, j]) + print("e2m1", out_a.e2m1_values[i, 8 * j : 8 * (j + 1)]) + print(backend_b) + print("sf", x_sf_b[i, j]) + print("e2m1", out_b.e2m1_values[i, 8 * j : 8 * (j + 1)]) + print("original") + print("x", x[i, 16 * j : 16 * (j + 1)]) + return + + if not torch.allclose(out_a.e2m1_values, out_b.e2m1_values): + mismatch_prop = ( + out_a.e2m1_values != out_b.e2m1_values + ).sum() / out_a.e2m1_values.numel() + print( + "Backends A and B have different e2m1 values! " + f"{mismatch_prop:.2%} mismatch", + ) + + [i, *_], [j, *_] = torch.where(out_a.e2m1_values != out_b.e2m1_values) + print(i, j) + print("normconst", out_a.amax) + print("sf", x_sf_a[i, j // 8]) + print(backend_a) + print("e2m1", out_a.e2m1_values[i, 8 * (j // 8) : 8 * (j // 8 + 1)]) + print(backend_b) + print("e2m1", out_b.e2m1_values[i, 8 * (j // 8) : 8 * (j // 8 + 1)]) + print("original") + print("x", x[i, 16 * (j // 8) : 16 * (j // 8 + 1)]) + return diff --git a/fouroversix/scripts/generate_kernels.py b/fouroversix/scripts/generate_kernels.py new file mode 100644 index 0000000..ccd286e --- /dev/null +++ b/fouroversix/scripts/generate_kernels.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +import argparse +import itertools +from dataclasses import dataclass +from pathlib import Path + +SRC_DTYPE_MAP = { + "fp16": "cutlass::half_t", + "bf16": "cutlass::bfloat16_t", +} + + +SM = [100] # Sm100 kernels support up to +IS_NVFP4 = ["false", "true"] +IS_TRANSPOSE = ["false", "true"] +IS_RHT = ["false", "true"] + + +def get_fp4_quant_template( + is_nvfp4: str, + is_rht: str, + is_transpose: str, + src_dtype: str, +) -> str: + if is_nvfp4 == "false": + function_str = "run_mxfp4_quant" + elif is_nvfp4 == "true": + function_str = "run_nvfp4_quant" + else: + msg = f"Invalid is_nvfp4: {is_nvfp4}" + raise ValueError(msg) + + if is_rht == "true": + function_str = f"{function_str}_rht" + + return f"""#include "fp4_quant_launch_template.h" +namespace fouroversix {{ + +template<> +void run_fp4_quant_<{src_dtype}, {is_nvfp4}, {is_rht}, {is_transpose}>(FP4_quant_params ¶ms, cudaStream_t stream) {{ + {function_str}<{src_dtype}, {is_transpose}>(params, stream); +}} + +}} // namespace fouroversix""" # noqa: E501 + + +@dataclass +class Kernel: + """Representation for a kernel that quantizes a tensor to FP4.""" + + sm: int + src_dtype: str + is_nvfp4: str + is_rht: str + is_transpose: str + op: str + + @property + def template(self) -> str: + """The kernel's template content.""" + + template_funcs = { + "fp4_quant": get_fp4_quant_template, + } + template_func = template_funcs[self.op] + return template_func( + is_transpose=self.is_transpose, + src_dtype=SRC_DTYPE_MAP[self.src_dtype], + is_nvfp4=self.is_nvfp4, + is_rht=self.is_rht, + ) + + @property + def filename(self) -> str: + """The filename for the kernel.""" + + fp4_format = "nvfp4" if self.is_nvfp4 == "true" else "mxfp4" + return ( + f"{self.op}_{self.src_dtype}_{fp4_format}_" + f"{'rht_' if self.is_rht == 'true' else ''}" + f"{'trans_' if self.is_transpose == 'true' else ''}sm{self.sm}.cu" + ) + + +def get_all_kernels() -> list[Kernel]: + for op in ["fp4_quant"]: + for src_dtype, is_nvfp4, is_rht, is_transpose, sm in itertools.product( + SRC_DTYPE_MAP.keys(), + IS_NVFP4, + IS_RHT, + IS_TRANSPOSE, + SM, + ): + yield Kernel( + sm=sm, + src_dtype=src_dtype, + is_rht=is_rht, + is_nvfp4=is_nvfp4, + is_transpose=is_transpose, + op=op, + ) + + +def write_kernel(kernel: Kernel, autogen_dir: Path) -> None: + prelude = """// Splitting the different transpose modes to different files to speed up compilation. +// This file is auto-generated. See "generate_kernels.py"\n""" # noqa: E501 + content = prelude + kernel.template + (autogen_dir / kernel.filename).write_text(content) + + +def main(output_dir: str | None) -> None: + if output_dir is None: + output_dir = ( + Path(__file__).parent.parent / "src" / "fouroversix" / "csrc" / "quantize" + ) + else: + output_dir = Path(output_dir) + + for kernel in get_all_kernels(): + write_kernel(kernel, output_dir) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + prog="generate_kernels", + description="Generate the flash_attention kernels template instantiations", + ) + parser.add_argument( + "-o", + "--output_dir", + required=False, + help="Where to generate the kernels will default to the current directory ", + ) + args = parser.parse_args() + main(args.output_dir) diff --git a/fouroversix/scripts/hadamard_code_gen.py b/fouroversix/scripts/hadamard_code_gen.py new file mode 100644 index 0000000..a9d3bd3 --- /dev/null +++ b/fouroversix/scripts/hadamard_code_gen.py @@ -0,0 +1,111 @@ +from pathlib import Path + +import numpy as np + +# had_16 = """ +# ++++++++++++++++ +# +-+-+-+-+-+-+-+- +# ++--++--++--++-- +# +--++--++--++--+ +# ++++----++++---- +# +-+--+-++-+--+-+ +# ++----++++----++ +# +--+-++-+--+-++- +# ++++++++-------- +# +-+-+-+--+-+-+-+ +# ++--++----++--++ +# +--++--+-++--++- +# ++++--------++++ +# +-+--+-+-+-++-+- +# ++----++--++++-- +# +--+-++--++-+--+ +# """ + +# transformerEngine style rht matrix +had_16 = """ ++++-+------+-+-- ++-++++-+-+-----+ +++-++-++--+--+++ ++---+++--+++--+- ++++--+++---++-++ ++-++--+--+--+++- +++-+-+----+-+--- ++------+-+++++-+ ++++-+---+++-+-++ ++-++++-++-+++++- +++-++-++++-++--- ++---+++-+---++-+ ++++--++++++--+-- ++-++--+-+-++---+ +++-+-+--++-+-+++ ++------++-----+- +""" + + +header = """ +/****************************************************************************** + * Copyright (c) 2023, Tri Dao. + * Adapted by Junxian Guo from https://github.com/Dao-AILab/fast-hadamard-transform/blob/master/csrc/code_gen.py + * Copyright (c) 2025, FourOverSix Team. + ******************************************************************************/ + +// This file is auto-generated. See "hadamard_code_gen.py"\n + +#pragma once + +""" + +template = """ +namespace fouroversix {{ + +__device__ __forceinline__ void hadamard_mult_thread_{N}(float x[{N}]) { + float out[{N}]; + {code} + #pragma unroll + for (int i = 0; i < {N}; i++) { x[i] = out[i]; } +} + +}} // namespace fouroversix + +""" + + +def string_to_array(string: str) -> np.ndarray: + # Convert strings of + and - to bool arrays + string = string.strip().replace("+", "1").replace("-", "-1").split() + return np.stack( + [ + np.fromstring(" ".join(string[i]), dtype=np.int32, sep=" ") + for i in range(len(string)) + ], + ) + + +def array_code_gen(arr: np.ndarray) -> str: + n = arr.shape[0] + if arr.shape[0] != arr.shape[1]: + msg = f"Hadamard matrix is not square: {arr.shape}" + raise ValueError(msg) + out = [ + f"out[{i}] = " + + " ".join([f"{'+' if arr[i, j] == 1 else '-'} x[{j}]" for j in range(n)]) + + ";" + for i in range(n) + ] + return template.replace("{N}", str(n)).replace("{code}", "\n ".join(out)) + + +def main() -> None: + output_dir = ( + Path(__file__).parent.parent + / "src" + / "fouroversix" + / "csrc" + / "include" + / "hadamard_transform_te.h" + ) + output_dir.write_text(header + array_code_gen(string_to_array(had_16))) + + +if __name__ == "__main__": + main() diff --git a/fouroversix/scripts/ptq/__init__.py b/fouroversix/scripts/ptq/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/fouroversix/scripts/ptq/__main__.py b/fouroversix/scripts/ptq/__main__.py new file mode 100644 index 0000000..b94a727 --- /dev/null +++ b/fouroversix/scripts/ptq/__main__.py @@ -0,0 +1,105 @@ +import warnings +from typing import Any + +import click +import modal +from fouroversix.utils import DataType, MatmulBackend, QuantizeBackend, ScaleRule + +from ..resources import app +from .coordinators import LocalEvaluationCoordinator, ModalEvaluationCoordinator +from .utils import EvaluationFramework, PTQMethod + + +@click.command() +@click.option( + "--activation-scale-rule", + "--a-scale-rule", + type=ScaleRule, + default=ScaleRule.mse, +) +@click.option("--detach", is_flag=True) +@click.option("--device", type=str, default="cuda") +@click.option("--dtype", type=DataType, default=DataType.nvfp4) +@click.option( + "--eval-framework", + "-f", + type=EvaluationFramework, + default=EvaluationFramework.lm_eval, +) +@click.option("--group-name", type=str, default=None) +@click.option("--limit", type=int, default=None) +@click.option("--matmul-backend", type=MatmulBackend, default=None) +@click.option("--max-length", type=int, default=None) +@click.option("--modal", is_flag=True) +@click.option("--modal-gpu", type=str) +@click.option("--model-name", "-m", type=str, multiple=True, required=True) +@click.option("--ptq-method", "-p", type=PTQMethod, multiple=True, required=True) +@click.option("--quantize-backend", type=QuantizeBackend, default=None) +@click.option("--task", "-t", type=str, multiple=True, default=["wikitext"]) +@click.option("--trust-remote-code", is_flag=True) +@click.option( + "--weight-scale-rule", + "--w-scale-rule", + type=ScaleRule, + default=ScaleRule.mse, +) +@click.option("--weight-scale-2d", "--w-scale-2d", is_flag=True) +def cli( + *, + detach: bool, + group_name: str | None, + modal_gpu: str, + **kwargs: dict[str, Any], +) -> None: + activation_scale_rule = kwargs.get("activation_scale_rule") + dtype = kwargs.get("dtype") + weight_scale_rule = kwargs.get("weight_scale_rule") + + model_names = kwargs.pop("model_name") + ptq_methods = kwargs.pop("ptq_method") + tasks = kwargs.pop("task") + use_modal = kwargs.pop("modal") + + # Expand shortcuts + if model_names[0] == "llamaqwen": + model_names = [ + "meta-llama/Llama-3.2-1B", + "meta-llama/Llama-3.1-8B", + "meta-llama/Llama-3.1-70B", + "Qwen/Qwen3-1.7B", + "Qwen/Qwen3-8B", + "Qwen/Qwen3-32B", + ] + + if isinstance(tasks, tuple): + tasks = list(tasks) + + if dtype == DataType.mxfp4 and ( + not activation_scale_rule.is_static() or not weight_scale_rule.is_static() + ): + msg = ( + "MXFP4 quantization only supports static scale rules. Setting " + "activation_scale_rule and weight_scale_rule to static_6..." + ) + warnings.warn(msg, stacklevel=1) + + kwargs["activation_scale_rule"] = ScaleRule.static_6 + kwargs["weight_scale_rule"] = ScaleRule.static_6 + + if use_modal: + with modal.enable_output(), app.run(detach=detach): + coordinator = ModalEvaluationCoordinator(group_name_str=group_name or "") + coordinator.start.remote( + model_names, + ptq_methods, + tasks, + modal_gpu=modal_gpu, + **kwargs, + ) + else: + coordinator = LocalEvaluationCoordinator(group_name) + coordinator.start(model_names, ptq_methods, tasks, **kwargs) + + +if __name__ == "__main__": + cli() diff --git a/fouroversix/scripts/ptq/coordinators/__init__.py b/fouroversix/scripts/ptq/coordinators/__init__.py new file mode 100644 index 0000000..67fb01e --- /dev/null +++ b/fouroversix/scripts/ptq/coordinators/__init__.py @@ -0,0 +1,4 @@ +from .local import LocalEvaluationCoordinator +from .modal import ModalEvaluationCoordinator + +__all__ = ["LocalEvaluationCoordinator", "ModalEvaluationCoordinator"] diff --git a/fouroversix/scripts/ptq/coordinators/base.py b/fouroversix/scripts/ptq/coordinators/base.py new file mode 100644 index 0000000..915fe54 --- /dev/null +++ b/fouroversix/scripts/ptq/coordinators/base.py @@ -0,0 +1,105 @@ +from abc import ABC, abstractmethod +from typing import Any + +from sqlalchemy import create_engine +from sqlalchemy.orm import Session, sessionmaker + +from ..experiment import Base, Experiment +from ..utils import PTQMethod + + +class BaseEvaluationCoordinator(ABC): + """Base class for evaluation coordinators.""" + + def get_session(self) -> Session: + """Get an SQLAlchemy session for the SQLite database.""" + engine = create_engine(f"sqlite:///{self.database_path.absolute().as_posix()}") + Base.metadata.create_all(engine) + return sessionmaker(bind=engine)() + + def get_tasks_to_evaluate( + self, + model_name: str, + ptq_method: PTQMethod, + tasks: list[str], + ) -> list[str]: + """ + Get the tasks that should be evaluated. If a group name is set, tasks will only + be evaluated if they have not yet been evaluated for this group name, model + name, PTQ method, and task. + """ + + if self.group_name is None: + return tasks + + session = self.get_session() + experiments = ( + session.query(Experiment) + .filter( + Experiment.group_name == self.group_name, + Experiment.model_name == model_name, + Experiment.ptq_method == ptq_method.value, + Experiment.task.in_(tasks), + ) + .all() + ) + + return [ + task + for task in tasks + if task not in [experiment.task for experiment in experiments] + ] + + @abstractmethod + def run_calibration_tasks( + self, + model_names: list[str], + ptq_methods: list[PTQMethod], + tasks: list[str], + **kwargs: dict[str, Any], + ) -> None: + """ + Run any tasks that should be used to calibrate models for a given PTQ method + and set of parameters before running evaluation. + """ + + def save_results( + self, + model_name: str, + ptq_method: PTQMethod, + kwargs: dict[str, Any], + results: list[tuple[str, str, float, dict[str, Any]]], + ) -> None: + """Save the results of a PTQ experiment to the SQLite database.""" + + session = self.get_session() + + for task, metric_name, metric_value, full_results in results: + experiment = Experiment( + group_name=self.group_name, + model_name=model_name, + task=task, + metric_name=metric_name, + metric_value=metric_value, + ptq_method=ptq_method.value, + activation_scale_rule=kwargs.get("activation_scale_rule"), + weight_scale_rule=kwargs.get("weight_scale_rule"), + smoothquant_alpha=kwargs.get("smoothquant_alpha"), + results=full_results, + ) + session.add(experiment) + + print(model_name, ptq_method, task) + print(full_results) + + session.commit() + + @abstractmethod + def start( + self, + model_names: list[str], + ptq_methods: list[PTQMethod], + tasks: list[str], + **kwargs: dict[str, Any], + ) -> None: + """Start the evaluation coordinator.""" diff --git a/fouroversix/scripts/ptq/coordinators/local.py b/fouroversix/scripts/ptq/coordinators/local.py new file mode 100644 index 0000000..59d7d7f --- /dev/null +++ b/fouroversix/scripts/ptq/coordinators/local.py @@ -0,0 +1,185 @@ +import itertools +import multiprocessing +from pathlib import Path +from typing import Any + +import torch + +from ..evaluators import get_evaluator +from ..utils import PTQMethod +from .base import BaseEvaluationCoordinator + +FOUROVERSIX_ROOT_DIR = Path(__file__).parent.parent.parent.parent + + +class LocalEvaluationCoordinator(BaseEvaluationCoordinator): + """Evaluation coordinator for running PTQ experiments locally.""" + + def __init__(self, group_name: str | None = None) -> None: + self.database_path = FOUROVERSIX_ROOT_DIR / "results.db" + self.group_name = group_name + + def evaluate( + self, + model_name: str, + ptq_method: PTQMethod, + **kwargs: dict[str, Any], + ) -> dict[str, Any]: + """Evaluate a model with a given PTQ method.""" + + evaluator_cls = get_evaluator(ptq_method) + + return evaluator_cls().evaluate( + model_name=model_name, + save_path=FOUROVERSIX_ROOT_DIR / "ptq", + **kwargs, + ) + + def run_calibration_tasks( + self, + model_names: list[str], + ptq_methods: list[PTQMethod], + tasks: list[str], + task_queue: multiprocessing.Queue, + result_queue: multiprocessing.Queue, + **kwargs: dict[str, Any], + ) -> None: + """ + Run any tasks that should be used to calibrate models for a given PTQ method + and set of parameters before running evaluation. + """ + + experiments = 0 + + for model_name, ptq_method in itertools.product(model_names, ptq_methods): + tasks_to_evaluate = self.get_tasks_to_evaluate( + model_name, + ptq_method, + tasks, + ) + + if len(tasks_to_evaluate) == 0: + continue + + evaluator_cls = get_evaluator(ptq_method) + + for calibration_task_kwargs in evaluator_cls.get_calibration_tasks( + model_name, + self.get_session(), + **kwargs, + ): + task_queue.put( + (model_name, ptq_method, {**kwargs, **calibration_task_kwargs}), + ) + experiments += 1 + + for _ in range(experiments): + self.save_results(*result_queue.get()) + + def start( + self, + model_names: list[str], + ptq_methods: list[PTQMethod], + tasks: list[str], + *, + device: str, + **kwargs: dict[str, Any], + ) -> None: + """Start the evaluation coordinator.""" + + multiprocessing.set_start_method("spawn", force=True) + + manager = multiprocessing.Manager() + task_queue = manager.Queue() + result_queue = manager.Queue() + + # Start one worker per GPU + num_workers = torch.cuda.device_count() if device == "cuda" else 1 + workers = [] + + for gpu_id in range(num_workers): + p = multiprocessing.Process( + target=self.worker, + args=( + f"cuda:{gpu_id}" if device == "cuda" else device, + task_queue, + result_queue, + ), + ) + p.start() + workers.append(p) + + # Run calibration tasks if necessary for each model and PTQ method + self.run_calibration_tasks( + model_names, + ptq_methods, + tasks, + task_queue, + result_queue, + **kwargs, + ) + + # Run evaluation tasks after models have been calibrated + experiments = 0 + + for model_name, ptq_method in itertools.product(model_names, ptq_methods): + tasks_to_evaluate = self.get_tasks_to_evaluate( + model_name, + ptq_method, + tasks, + ) + + if len(tasks_to_evaluate) == 0: + continue + + evaluator_cls = get_evaluator(ptq_method) + + calibrated_kwargs = evaluator_cls.get_calibrated_kwargs( + model_name, + self.get_session(), + **kwargs, + ) + + task_queue.put( + ( + model_name, + ptq_method, + {**kwargs, "tasks": tasks_to_evaluate, **calibrated_kwargs}, + ), + ) + experiments += 1 + + # Send shutdown signals (one per worker) + for _ in range(num_workers): + task_queue.put(None) + + # Collect results + for _ in range(experiments): + self.save_results(*result_queue.get()) + + for p in workers: + p.join() + + def worker( + self, + device: str, + task_queue: multiprocessing.Queue, + result_queue: multiprocessing.Queue, + ) -> None: + """Worker process for running PTQ experiments locally.""" + + while True: + worker_task = task_queue.get() + + if worker_task is None: + break + + model_name, ptq_method, kwargs = worker_task + + results = self.evaluate( + model_name, + ptq_method, + **{**kwargs, "device": device}, + ) + + result_queue.put((model_name, ptq_method, kwargs, results)) diff --git a/fouroversix/scripts/ptq/coordinators/modal.py b/fouroversix/scripts/ptq/coordinators/modal.py new file mode 100644 index 0000000..96261d6 --- /dev/null +++ b/fouroversix/scripts/ptq/coordinators/modal.py @@ -0,0 +1,152 @@ +import itertools +from pathlib import Path +from typing import Any + +import modal + +from ...resources import FOUROVERSIX_CACHE_PATH, app, cache_volume, get_image +from ..evaluators import get_evaluator +from ..utils import PTQMethod +from .base import BaseEvaluationCoordinator + + +@app.cls( + image=get_image(), + timeout=24 * 60 * 60, + nonpreemptible=True, + volumes={FOUROVERSIX_CACHE_PATH: cache_volume}, +) +class ModalEvaluationCoordinator(BaseEvaluationCoordinator): + """Evaluation coordinator for running PTQ experiments on Modal.""" + + group_name_str: str = modal.parameter() + + @property + def database_path(self) -> Path: + """Path to the SQLite database where experiment results are stored.""" + return FOUROVERSIX_CACHE_PATH / "results.db" + + @property + def group_name(self) -> str | None: + """ + The name of the group experiments are being run in. If this is not None and an + experiment with this group name and matching parameters has already been run, + the experiment will not be run again. + """ + + # Modal doesn't allow None parameters in modal.parameter() + return self.group_name_str if self.group_name_str != "" else None + + def run_calibration_tasks( + self, + model_names: list[str], + ptq_methods: list[PTQMethod], + tasks: list[str], + modal_gpu: str, + **kwargs: dict[str, Any], + ) -> None: + """ + Run any tasks that should be used to calibrate models for a given PTQ method + and set of parameters before running evaluation. + """ + + function_calls_with_inputs = [] + + for model_name, ptq_method in itertools.product(model_names, ptq_methods): + tasks_to_evaluate = self.get_tasks_to_evaluate( + model_name, + ptq_method, + tasks, + ) + + if len(tasks_to_evaluate) == 0: + continue + + evaluator_cls = get_evaluator(ptq_method).with_options(gpu=modal_gpu) + + function_calls_with_inputs.extend( + [ + ( + model_name, + ptq_method, + {**kwargs, **calibration_task_kwargs}, + evaluator_cls().evaluate_on_modal.spawn( + model_name=model_name, + save_path=FOUROVERSIX_CACHE_PATH / "ptq", + **{ + **kwargs, + "tasks": tasks_to_evaluate, + **calibration_task_kwargs, + }, + ), + ) + for calibration_task_kwargs in evaluator_cls.get_calibration_tasks( + model_name, + self.get_session(), + **kwargs, + ) + ], + ) + + results = modal.FunctionCall.gather( + *[function_call for _, _, _, function_call in function_calls_with_inputs], + ) + + for (model_name, ptq_method, function_call_kwargs, _), result in zip( + function_calls_with_inputs, + results, + strict=True, + ): + self.save_results(model_name, ptq_method, function_call_kwargs, result) + + @modal.method() + def start( + self, + model_names: list[str], + ptq_methods: list[PTQMethod], + tasks: list[str], + modal_gpu: str, + **kwargs: dict[str, Any], + ) -> None: + """Start the evaluation coordinator.""" + + self.run_calibration_tasks(model_names, ptq_methods, tasks, modal_gpu, **kwargs) + + models_and_ptq_methods = list(itertools.product(model_names, ptq_methods)) + function_calls = [] + + for model_name, ptq_method in models_and_ptq_methods: + tasks_to_evaluate = self.get_tasks_to_evaluate( + model_name, + ptq_method, + tasks, + ) + + if len(tasks_to_evaluate) == 0: + continue + + evaluator_cls = get_evaluator(ptq_method).with_options(gpu=modal_gpu) + + calibrated_kwargs = evaluator_cls.get_calibrated_kwargs( + model_name, + self.get_session(), + **kwargs, + ) + + function_calls.append( + evaluator_cls().evaluate_on_modal.spawn( + model_name=model_name, + tasks=tasks_to_evaluate, + save_path=FOUROVERSIX_CACHE_PATH / "ptq", + **{**kwargs, **calibrated_kwargs}, + ), + ) + + all_results = modal.FunctionCall.gather(*function_calls) + + for (model_name, ptq_method), results in zip( + models_and_ptq_methods, + all_results, + strict=True, + ): + self.save_results(model_name, ptq_method, kwargs, results) diff --git a/fouroversix/scripts/ptq/evaluators/__init__.py b/fouroversix/scripts/ptq/evaluators/__init__.py new file mode 100644 index 0000000..b13f711 --- /dev/null +++ b/fouroversix/scripts/ptq/evaluators/__init__.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from ..utils import PTQMethod +from .awq import AWQEvaluator +from .gptq import GPTQEvaluator +from .high_precision import HighPrecisionEvaluator +from .rtn import RTNEvaluator +from .smoothquant import SmoothQuantEvaluator +from .spinquant import SpinQuantEvaluator + +if TYPE_CHECKING: + from .evaluator import PTQEvaluator + + +def get_evaluator(ptq_method: PTQMethod) -> type[PTQEvaluator]: + """Get the evaluator class for the given PTQ method.""" + + if ptq_method == PTQMethod.awq: + return AWQEvaluator + if ptq_method == PTQMethod.gptq: + return GPTQEvaluator + if ptq_method == PTQMethod.high_precision: + return HighPrecisionEvaluator + if ptq_method == PTQMethod.rtn: + return RTNEvaluator + if ptq_method == PTQMethod.smoothquant: + return SmoothQuantEvaluator + if ptq_method == PTQMethod.spinquant: + return SpinQuantEvaluator + + msg = f"Unsupported PTQ method: {ptq_method}" + raise ValueError(msg) diff --git a/fouroversix/scripts/ptq/evaluators/awq.py b/fouroversix/scripts/ptq/evaluators/awq.py new file mode 100644 index 0000000..0fb7b60 --- /dev/null +++ b/fouroversix/scripts/ptq/evaluators/awq.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import torch +import torch.nn as nn +import torch.nn.functional as F +from fouroversix import FourOverSixLinear, ModelQuantizationConfig +from fouroversix.model.quantize import QuantizedModule + +from ...resources import ( + FOUROVERSIX_CACHE_PATH, + Dependency, + app, + cache_volume, + get_image, + hf_secret, +) +from .rtn import RTNEvaluatorImpl + +if TYPE_CHECKING: + from transformers import AutoModelForCausalLM + +awq_img = get_image(dependencies=[Dependency.fouroversix, Dependency.awq]) + + +class FourOverSixLinearForAWQ(FourOverSixLinear): + """ + Drop-in replacement for `FourOverSixLinear` that quantizes the weights and + activations during AWQ calibration. + """ + + def __init__(self, *args: list[Any], **kwargs: dict[str, Any]) -> None: + super().__init__(*args, **kwargs) + self.config.keep_master_weights = True + self.high_precision = False + + def apply_ptq(self) -> None: + """ + Override the parent method to do nothing, since we need the high-precision + weight when calibrating with AWQ. + """ + + def forward(self, input: torch.Tensor) -> torch.Tensor: + """ + Forward pass that can optionally be run in high precision. This is used to + calculate the high-precision output to compare to during the auto-scale process + in AWQ calibration. + """ + + return ( + F.linear(input, self.weight, self.bias) + if self.high_precision + else super().forward(input) + ) + + +@app.cls( + image=awq_img, + gpu="B200", + secrets=[hf_secret], + timeout=24 * 60 * 60, + volumes={FOUROVERSIX_CACHE_PATH.as_posix(): cache_volume}, +) +class AWQEvaluator(RTNEvaluatorImpl): + """Evaluate a model using AWQ.""" + + def quantize_model( + self, + model_name: str, + *, + device: str, + save_path: Path, + quantization_config: ModelQuantizationConfig, + trust_remote_code: bool, + ) -> AutoModelForCausalLM: + """Quantize a model using AWQ.""" + + import torch + from awq.quantize.pre_quant import apply_awq, run_awq + from fouroversix import quantize_model + from transformers import AutoModelForCausalLM, AutoTokenizer + + # Replace FourOverSixLinear with FourOverSixLinearForAWQ + QuantizedModule.register( + nn.Linear, + replace_existing_modules_in_registry=True, + )(FourOverSixLinearForAWQ) + + save_path = save_path / "awq" / model_name / quantization_config.__hash__() + + if not save_path.exists(): + model = AutoModelForCausalLM.from_pretrained( + model_name, + device_map=device, + trust_remote_code=trust_remote_code, + ).eval() + + quantize_model(model, quantization_config) + + enc = AutoTokenizer.from_pretrained( + model_name, + use_fast=False, + trust_remote_code=trust_remote_code, + ) + + awq_results = run_awq( + model, + enc, + w_bit=16, + q_config={"q_group_size": -1, "zero_point": False}, + n_samples=128, + seqlen=512, + calib_data="wikitext", + ) + + Path(save_path).parent.mkdir(parents=True, exist_ok=True) + torch.save(awq_results, save_path) + + model = AutoModelForCausalLM.from_pretrained( + model_name, + device_map=device, + trust_remote_code=trust_remote_code, + ) + + # Apply AWQ + awq_results = torch.load(save_path, map_location="cuda") + apply_awq(model, awq_results) + + # Quantize the model + quantize_model(model, quantization_config) + + return model.to(device) diff --git a/fouroversix/scripts/ptq/evaluators/evaluator.py b/fouroversix/scripts/ptq/evaluators/evaluator.py new file mode 100644 index 0000000..5559992 --- /dev/null +++ b/fouroversix/scripts/ptq/evaluators/evaluator.py @@ -0,0 +1,204 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from contextlib import nullcontext +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import modal +import torch +from fouroversix import ( + DataType, + MatmulBackend, + ModelQuantizationConfig, + QuantizeBackend, + ScaleRule, +) + +from ..utils import EvaluationFramework + +if TYPE_CHECKING: + from sqlalchemy.orm import Session + +from transformers import AutoConfig, AutoModelForCausalLM + + +class PTQEvaluator(ABC): + """Base class for post-training quantization evaluators.""" + + @classmethod + def get_calibration_tasks( + cls, + model_name: str, # noqa: ARG003 + session: Session, # noqa: ARG003 + **kwargs: dict[str, Any], # noqa: ARG003 + ) -> list[dict[str, Any]]: + """ + Get the kwargs for tasks that should be used to calibrate the given model for + this PTQ method before running evaluation. + """ + return [] + + @classmethod + def get_calibrated_kwargs( + cls, + model_name: str, # noqa: ARG003 + session: Session, # noqa: ARG003 + **kwargs: dict[str, Any], # noqa: ARG003 + ) -> dict[str, Any]: + """ + Get the calibrated kwargs for the given model and scale rules. If this model + has not yet been calibrated with these scale rules, an error will be raised. + """ + return {} + + @abstractmethod + def quantize_model(self, **kwargs: dict[str, Any]) -> AutoModelForCausalLM: + """Quantize a model.""" + + def evaluate( + self, + model_name: str, + *, + device: str, + dtype: str, + eval_framework: EvaluationFramework, + limit: int | None, + max_length: int | None, + tasks: list[str], + trust_remote_code: bool = False, + disable_inference_mode: bool = False, + matmul_backend: MatmulBackend | None = None, + quantize_backend: QuantizeBackend | None = None, + weight_scale_2d: bool = False, + activation_scale_rule: ScaleRule | None = None, + weight_scale_rule: ScaleRule | None = None, + save_path: Path | None = None, + **kwargs: dict[str, Any], + ) -> dict[str, Any]: + """Evaluate a quantized model with lm-eval.""" + + inference_context = ( + nullcontext() if disable_inference_mode else torch.inference_mode() + ) + + with inference_context: + model_config = AutoConfig.from_pretrained(model_name) + quantization_config = ModelQuantizationConfig( + activation_scale_rule=activation_scale_rule, + dtype=dtype, + matmul_backend=matmul_backend, + output_dtype=DataType( + ( + str(model_config.dtype).replace("torch.", "") + if model_config.dtype is not None + else "bfloat16" + ), + ), + quantize_backend=quantize_backend, + weight_scale_2d=weight_scale_2d, + weight_scale_rule=weight_scale_rule, + ) + + model = self.quantize_model( + model_name=model_name, + device=device, + save_path=save_path, + quantization_config=quantization_config, + trust_remote_code=trust_remote_code, + **kwargs, + ) + + if eval_framework == EvaluationFramework.lm_eval: + from lm_eval import evaluator + from lm_eval.models.huggingface import HFLM + from lm_eval.tasks import TaskManager + + full_results = evaluator.simple_evaluate( + model=HFLM( + pretrained=model, + device=device, + max_length=max_length, + ), + tasks=tasks, + device=device, + limit=limit, + task_manager=TaskManager( + include_path=( + Path(__file__).parent.parent / "tasks" + ).as_posix(), + ), + ) + + results = [] + + for task in full_results["results"]: + result = full_results["results"][task] + + if "acc_norm,none" in result: + metric_name = "acc_norm,none" + elif "acc,none" in result: + metric_name = "acc,none" + elif "word_perplexity,none" in result: + metric_name = "word_perplexity,none" + else: + metric_name = None + + results.append( + ( + task, + metric_name, + result.get(metric_name), + full_results["results"][task], + ), + ) + + elif eval_framework == EvaluationFramework.inspect_ai: + import inspect_ai + from inspect_ai.model import Model + from inspect_ai.model._generate_config import GenerateConfig + + from .utils import local_hf + + config = GenerateConfig() + full_results = inspect_ai.eval( + tasks=tasks, + model=Model(local_hf(model_name, model, config), config, None), + limit=limit, + log_dir=(save_path / "inspect_ai_logs").as_posix(), + display="none", + ) + + results = [] + + for log in full_results: + metrics = { + k: v.value + for score in log.results.scores + for k, v in score.metrics.items() + } + + metric_name = "accuracy" if "accuracy" in metrics else None + + results.append( + ( + log.eval.task, + metric_name, + metrics.get(metric_name), + metrics, + ), + ) + + del model + torch.cuda.empty_cache() + + return results + + @modal.method() + def evaluate_on_modal( + self, + *args: list[Any], + **kwargs: dict[str, Any], + ) -> dict[str, Any]: + """Evaluate a quantized model on Modal.""" + return self.evaluate(*args, **kwargs) diff --git a/fouroversix/scripts/ptq/evaluators/gptq.py b/fouroversix/scripts/ptq/evaluators/gptq.py new file mode 100644 index 0000000..00e1e4a --- /dev/null +++ b/fouroversix/scripts/ptq/evaluators/gptq.py @@ -0,0 +1,99 @@ +import sys +from pathlib import Path +from typing import TYPE_CHECKING + +import fouroversix +from fouroversix import ModelQuantizationConfig + +from ...resources import ( + FOUROVERSIX_CACHE_PATH, + Dependency, + app, + cache_volume, + get_image, + hf_secret, +) +from .evaluator import PTQEvaluator + +if TYPE_CHECKING: + from transformers import AutoModelForCausalLM + +CALIBRATION_DATASET = "wikitext" + +gptq_img = get_image( + dependencies=[ + Dependency.fast_hadamard_transform, + Dependency.qutlass, + Dependency.fp_quant, + Dependency.fouroversix, + ], +) + + +@app.cls( + image=gptq_img, + gpu="B200", + secrets=[hf_secret], + timeout=24 * 60 * 60, + volumes={FOUROVERSIX_CACHE_PATH: cache_volume}, +) +class GPTQEvaluator(PTQEvaluator): + """Evaluate a model after quantizing it with GPTQ.""" + + def quantize_model( + self, + model_name: str, + *, + device: str, + save_path: Path, + quantization_config: ModelQuantizationConfig, + trust_remote_code: bool, + ) -> "AutoModelForCausalLM": + """Quantize a model with GPTQ.""" + + sys.path.extend( + [ + ( + Path(fouroversix.__file__).parent.parent.parent + / "third_party" + / "fp-quant" + ).as_posix(), + ], + ) + + from model_quant import main + from transformers import AutoModelForCausalLM + + save_path = save_path / "gptq" / model_name / quantization_config.__hash__() + + if not save_path.exists(): + sys.argv = [ + sys.argv[0], + "--model_name_or_path", + model_name, + "--dataset_name_or_path", + CALIBRATION_DATASET, + "--w_bits", + "4", + "--a_bits", + "4", + "--export_quantized_model", + "realquant", + "--format", + "nvfp", + "--gptq", + "--save_path", + save_path.as_posix(), + "--a_scale_rule", + quantization_config.activation_scale_rule.value, + "--w_scale_rule", + quantization_config.weight_scale_rule.value, + ] + + main() + + return AutoModelForCausalLM.from_pretrained( + save_path, + device_map=device, + trust_remote_code=trust_remote_code, + ) diff --git a/fouroversix/scripts/ptq/evaluators/high_precision.py b/fouroversix/scripts/ptq/evaluators/high_precision.py new file mode 100644 index 0000000..9855290 --- /dev/null +++ b/fouroversix/scripts/ptq/evaluators/high_precision.py @@ -0,0 +1,45 @@ +from pathlib import Path + +from fouroversix import ModelQuantizationConfig + +from ...resources import ( + FOUROVERSIX_CACHE_PATH, + app, + cache_volume, + get_image, + hf_secret, +) +from .evaluator import PTQEvaluator + +hp_img = get_image() + +with hp_img.imports(): + from transformers import AutoModelForCausalLM + + +@app.cls( + image=hp_img, + gpu="B200", + secrets=[hf_secret], + timeout=24 * 60 * 60, + volumes={FOUROVERSIX_CACHE_PATH.as_posix(): cache_volume}, +) +class HighPrecisionEvaluator(PTQEvaluator): + """Evaluate a model while keeping it in high precision.""" + + def quantize_model( + self, + model_name: str, + *, + device: str, + save_path: Path, # noqa: ARG002 + quantization_config: ModelQuantizationConfig, # noqa: ARG002 + trust_remote_code: bool = False, + ) -> "AutoModelForCausalLM": + """Return a model without any quantization.""" + + return AutoModelForCausalLM.from_pretrained( + model_name, + device_map=device, + trust_remote_code=trust_remote_code, + ) diff --git a/fouroversix/scripts/ptq/evaluators/rtn.py b/fouroversix/scripts/ptq/evaluators/rtn.py new file mode 100644 index 0000000..8e23abe --- /dev/null +++ b/fouroversix/scripts/ptq/evaluators/rtn.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from ...resources import ( + FOUROVERSIX_CACHE_PATH, + app, + cache_volume, + get_image, + hf_secret, +) +from .evaluator import PTQEvaluator + +if TYPE_CHECKING: + from pathlib import Path + + from fouroversix import ModelQuantizationConfig + from transformers import AutoModelForCausalLM + + +rtn_img = get_image() + +with rtn_img.imports(): + from transformers import AutoConfig, AutoModelForCausalLM + + try: + from transformers import FourOverSixConfig as HFFourOverSixConfig + except ImportError: + HFFourOverSixConfig = None + + +class RTNEvaluatorImpl(PTQEvaluator): + """Evaluate a model using round-to-nearest quantization.""" + + def quantize_model( + self, + model_name: str, + *, + device: str, + save_path: Path, + quantization_config: ModelQuantizationConfig, + trust_remote_code: bool = False, + ) -> AutoModelForCausalLM: + """Quantize a model using round-to-nearest quantization.""" + + model_save_path = ( + save_path / "rtn" / model_name / quantization_config.__hash__() + ) + + if not model_save_path.exists(): + model_config = AutoConfig.from_pretrained(model_name) + + hf_quantization_config = HFFourOverSixConfig( + activation_scale_rule=quantization_config.activation_scale_rule, + dtype=quantization_config.dtype, + matmul_backend=quantization_config.matmul_backend, + output_dtype=quantization_config.output_dtype, + quantize_backend=quantization_config.quantize_backend, + weight_scale_2d=quantization_config.weight_scale_2d, + weight_scale_rule=quantization_config.weight_scale_rule, + ) + + save_kwargs = {} + if hasattr(model_config, "quantization_config"): + hf_quantization_config.pre_quantized_model_config_type = str( + type(model_config), + ) + save_kwargs["save_original_format"] = False + delattr(model_config, "quantization_config") + + model = AutoModelForCausalLM.from_pretrained( + model_name, + device_map=device, + config=model_config, + quantization_config=hf_quantization_config, + trust_remote_code=trust_remote_code, + ) + + if hasattr(hf_quantization_config, "pre_quantized_model_config_type"): + delattr(hf_quantization_config, "pre_quantized_model_config_type") + + model.save_pretrained(model_save_path, **save_kwargs) + else: + model = AutoModelForCausalLM.from_pretrained( + model_save_path, + device_map=device, + trust_remote_code=trust_remote_code, + ) + + # Fix for Inspect AI + model.name_or_path = model_name + + return model + + +@app.cls( + image=rtn_img, + cpu=4, + memory=8 * 1024, + gpu="B200", + secrets=[hf_secret], + timeout=24 * 60 * 60, + volumes={FOUROVERSIX_CACHE_PATH.as_posix(): cache_volume}, +) +class RTNEvaluator(RTNEvaluatorImpl): + """Evaluate a model using round-to-nearest quantization.""" diff --git a/fouroversix/scripts/ptq/evaluators/smoothquant.py b/fouroversix/scripts/ptq/evaluators/smoothquant.py new file mode 100644 index 0000000..05b6d02 --- /dev/null +++ b/fouroversix/scripts/ptq/evaluators/smoothquant.py @@ -0,0 +1,238 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from fouroversix import ModelQuantizationConfig, ScaleRule + +from ...resources import FOUROVERSIX_CACHE_PATH, app, cache_volume, hf_secret +from ..experiment import Experiment +from ..utils import PTQMethod +from .rtn import RTNEvaluatorImpl, rtn_img + +if TYPE_CHECKING: + from pathlib import Path + + from sqlalchemy.orm import Session + + +with rtn_img.imports(): + import torch + import torch.nn as nn + from fouroversix import ( + FourOverSixLinear, + QuantizedModule, + fp4_matmul, + quantize_model, + ) + from transformers import AutoModelForCausalLM + + +ALPHA_CANDIDATES = [x / 10 for x in range(11)] +WIKITEXT_TRAIN = "wikitext_train" + + +class FourOverSixLinearWithSmoothing(FourOverSixLinear): + """ + Drop-in replacement for `FourOverSixLinear` that implements SmoothQuant-style + scaling. + """ + + def __init__( + self, + *args: list[Any], + smoothquant_alpha: float, + **kwargs: dict[str, Any], + ) -> None: + super().__init__(*args, **kwargs) + self.smoothquant_alpha = smoothquant_alpha + + def apply_ptq(self) -> None: + """ + Override the parent method to do nothing, since we need the high-precision + weight when doing PTQ with SmoothQuant. + """ + + def forward(self, input: torch.Tensor) -> torch.Tensor: + """Forward pass with SmoothQuant-style scaling.""" + + out = torch.empty( + *input.shape[:-1], + self.weight.shape[0], + device=input.device, + dtype=self.config.output_dtype.torch_dtype(), + ) + + fprop_activation_config = self.config.get_activation_config() + fprop_weight_config = self.config.get_weight_config( + block_scale_2d=self.config.weight_scale_2d, + ) + + for i in range(input.shape[0]): + s = (input[i].abs().max(dim=0).values ** self.smoothquant_alpha) / ( + self.weight.abs().max(dim=0).values ** (1 - self.smoothquant_alpha) + ) + + out[i] = fp4_matmul( + input[i] / s[None, :], + self.weight * s[None, :], + out_dtype=self.config.output_dtype, + input_config=fprop_activation_config, + other_config=fprop_weight_config, + ) + + if self.bias is not None: + out = out + self.bias + + return out + + +@app.cls( + image=rtn_img, + gpu="B200", + secrets=[hf_secret], + timeout=24 * 60 * 60, + volumes={FOUROVERSIX_CACHE_PATH.as_posix(): cache_volume}, +) +class SmoothQuantEvaluator(RTNEvaluatorImpl): + """Evaluate a model using SmoothQuant.""" + + @classmethod + def get_calibration_tasks( + cls, + model_name: str, + session: Session, + **kwargs: dict[str, Any], + ) -> list[dict[str, Any]]: + """ + Get the kwargs for tasks that should be used to calibrate the given model for + this PTQ method before running evaluation. + """ + + smoothquant_alpha = get_smoothquant_alpha( + model_name, + kwargs.get("activation_scale_rule"), + kwargs.get("weight_scale_rule"), + session, + ) + + calibration_experiments = get_calibration_experiments( + model_name, + kwargs.get("activation_scale_rule"), + kwargs.get("weight_scale_rule"), + session, + ) + + if smoothquant_alpha is None: + return [ + { + "smoothquant_alpha": candidate_alpha, + "tasks": [WIKITEXT_TRAIN], + } + for candidate_alpha in ALPHA_CANDIDATES + if not any( + experiment.smoothquant_alpha == candidate_alpha + for experiment in calibration_experiments + ) + ] + + return [] + + @classmethod + def get_calibrated_kwargs( + cls, + model_name: str, + session: Session, + **kwargs: dict[str, Any], + ) -> dict[str, Any]: + """ + Get the calibrated kwargs for the given model and scale rules. If this model + has not yet been calibrated with these scale rules, an error will be raised. + """ + + smoothquant_alpha = get_smoothquant_alpha( + model_name, + kwargs.get("activation_scale_rule"), + kwargs.get("weight_scale_rule"), + session, + ) + + if smoothquant_alpha is None: + msg = ( + "SmoothQuant has not been calibrated for this combination of model and " + "scale rules" + ) + raise ValueError(msg) + + return {"smoothquant_alpha": smoothquant_alpha} + + def quantize_model( + self, + model_name: str, + *, + device: str, + save_path: Path, # noqa: ARG002 + smoothquant_alpha: float, + quantization_config: ModelQuantizationConfig, + trust_remote_code: bool, + ) -> AutoModelForCausalLM: + """Quantize a model using SmoothQuant.""" + + # Replace FourOverSixLinear with FourOverSixLinearWithSmoothing + QuantizedModule.register( + nn.Linear, + replace_existing_modules_in_registry=True, + )(FourOverSixLinearWithSmoothing) + + model = AutoModelForCausalLM.from_pretrained( + model_name, + device_map=device, + trust_remote_code=trust_remote_code, + ) + + quantize_model(model, quantization_config, smoothquant_alpha=smoothquant_alpha) + return model + + +def get_calibration_experiments( + model_name: str, + activation_scale_rule: ScaleRule, + weight_scale_rule: ScaleRule, + db_session: Session, +) -> list[Experiment]: + return ( + db_session.query(Experiment) + .filter( + Experiment.ptq_method == PTQMethod.smoothquant.value, + Experiment.task == WIKITEXT_TRAIN, + Experiment.model_name == model_name, + Experiment.activation_scale_rule == activation_scale_rule.value, + Experiment.weight_scale_rule == weight_scale_rule.value, + Experiment.smoothquant_alpha.isnot(None), + ) + .all() + ) + + +def get_smoothquant_alpha( + model_name: str, + activation_scale_rule: ScaleRule, + weight_scale_rule: ScaleRule, + session: Session, +) -> float | None: + calibration_experiments = get_calibration_experiments( + model_name, + activation_scale_rule, + weight_scale_rule, + session, + ) + + if not all( + any( + experiment.smoothquant_alpha == alpha + for experiment in calibration_experiments + ) + for alpha in ALPHA_CANDIDATES + ): + return None + + return min(calibration_experiments, key=lambda x: x.metric_value).smoothquant_alpha diff --git a/fouroversix/scripts/ptq/evaluators/spinquant.py b/fouroversix/scripts/ptq/evaluators/spinquant.py new file mode 100644 index 0000000..7a22a04 --- /dev/null +++ b/fouroversix/scripts/ptq/evaluators/spinquant.py @@ -0,0 +1,240 @@ +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import fouroversix +import modal +from fouroversix import ModelQuantizationConfig + +from ...resources import ( + FOUROVERSIX_CACHE_PATH, + Dependency, + app, + cache_volume, + get_image, + hf_secret, +) +from ..utils import get_model_size +from .evaluator import PTQEvaluator + +if TYPE_CHECKING: + from transformers import AutoModelForCausalLM + +spinquant_img = get_image( + dependencies=[ + Dependency.fast_hadamard_transform, + Dependency.fouroversix, + Dependency.spinquant, + ], + extra_pip_dependencies=["transformers<5.0"], +) + + +MIN_MODEL_SIZE_FOR_8xB200 = 32 +SPINQUANT_STEPS = 100 + +SPINQUANT_ARGS = [ + "--model_max_length", + "8192", + "--fp16", + "False", + "--bf16", + "True", + "--w_bits", + "4", + "--a_bits", + "4", + "--k_bits", + "16", + "--v_bits", + "16", +] + + +@app.cls( + image=spinquant_img, + timeout=24 * 60 * 60, + secrets=[hf_secret], + volumes={FOUROVERSIX_CACHE_PATH.as_posix(): cache_volume}, +) +class SpinQuantOptimizer: + """Optimize a model with SpinQuant.""" + + def optimize( + self, + model_name: str, + *, + quantization_config: ModelQuantizationConfig, + spinquant_save_path: str, + spinquant_steps: int, + ) -> None: + """Optimize a model with SpinQuant.""" + + subprocess.run( + [ + "torchrun", + "--nnodes=1", + "--nproc_per_node=auto", + ( + Path(fouroversix.__file__).parent.parent.parent + / "third_party" + / "spinquant" + / "optimize_rotation.py" + ).as_posix(), + "--input_model", + model_name, + "--output_dir", + spinquant_save_path, + "--output_rotation_path", + spinquant_save_path, + "--log_on_each_node", + "False", + "--per_device_train_batch_size", + "1", + "--logging_steps", + "1", + "--learning_rate", + "1.5", + "--weight_decay", + "0.", + "--lr_scheduler_type", + "cosine", + "--gradient_checkpointing", + "True", + "--save_safetensors", + "False", + "--max_steps", + str(spinquant_steps), + "--activation_scale_rule", + quantization_config.activation_scale_rule.value, + "--weight_scale_rule", + quantization_config.weight_scale_rule.value, + *SPINQUANT_ARGS, + ], + check=True, + ) + + cache_volume.commit() + + @modal.method() + def optimize_on_modal( + self, + *args: list[Any], + **kwargs: dict[str, Any], + ) -> None: + """Optimize a model with SpinQuant on Modal.""" + return self.optimize(*args, **kwargs) + + +@app.cls( + image=spinquant_img, + timeout=24 * 60 * 60, + secrets=[hf_secret], + gpu="B200", + volumes={FOUROVERSIX_CACHE_PATH.as_posix(): cache_volume}, +) +class SpinQuantEvaluator(PTQEvaluator): + """Evaluate a quantized model with SpinQuant.""" + + def quantize_model( + self, + model_name: str, + *, + device: str, + save_path: Path, + quantization_config: ModelQuantizationConfig, + trust_remote_code: bool, + ) -> AutoModelForCausalLM: + """Export a quantized model with SpinQuant.""" + + import fouroversix + + sys.path.append( + ( + Path(fouroversix.__file__).parent.parent.parent + / "third_party" + / "spinquant" + ).as_posix(), + ) + + from eval_utils.main import ptq_model + from transformers import AutoConfig, AutoModelForCausalLM + from utils.process_args import process_args_ptq + + save_path = ( + save_path + / "spinquant" + / ( + f"{model_name}-{quantization_config.activation_scale_rule.value}" + f"-{quantization_config.weight_scale_rule.value}" + ) + ) + + if not (save_path / "R.bin").exists(): + model_is_large = get_model_size(model_name) >= MIN_MODEL_SIZE_FOR_8xB200 + + if model_is_large: + msg = ( + "Automatic SpinQuant optimization is not supported for large " + "models. Please optimize the model manually." + ) + raise RuntimeError(msg) + + SpinQuantOptimizer().optimize( + model_name, + quantization_config=quantization_config, + spinquant_save_path=save_path.as_posix(), + spinquant_steps=SPINQUANT_STEPS, + ) + + sys.argv = [ + sys.argv[0], + "--input_model", + model_name, + "--do_train", + "False", + "--do_eval", + "True", + "--per_device_eval_batch_size", + "4", + "--rotate", + "--optimized_rotation_path", + (save_path / "R.bin").as_posix(), + "--activation_scale_rule", + quantization_config.activation_scale_rule.value, + "--weight_scale_rule", + quantization_config.weight_scale_rule.value, + *SPINQUANT_ARGS, + ] + + config = AutoConfig.from_pretrained(model_name) + + # Llama v3.2 specific: Spinquant is not compatiable with tie_word_embeddings, + # clone lm_head from embed_tokens + process_word_embeddings = False + if config.tie_word_embeddings: + config.tie_word_embeddings = False + process_word_embeddings = True + + model = AutoModelForCausalLM.from_pretrained( + model_name, + config=config, + device_map=device, + trust_remote_code=trust_remote_code, + ) + + if process_word_embeddings: + model.lm_head.weight.data = model.model.embed_tokens.weight.data.clone() + + model.to(device) + model_args, _, ptq_args = process_args_ptq() + + cache_volume.reload() + + model = ptq_model(ptq_args, model, model_args) + model.to(device) + + return model diff --git a/fouroversix/scripts/ptq/evaluators/utils.py b/fouroversix/scripts/ptq/evaluators/utils.py new file mode 100644 index 0000000..8ffdbdf --- /dev/null +++ b/fouroversix/scripts/ptq/evaluators/utils.py @@ -0,0 +1,100 @@ +from typing import Any + +import torch +from inspect_ai.model import modelapi +from inspect_ai.model._generate_config import GenerateConfig +from inspect_ai.model._providers.hf import HuggingFaceAPI +from transformers import AutoModelForCausalLM, AutoTokenizer + + +def set_random_seeds(seed: int | None = None) -> None: + import os + + import numpy as np + from transformers import set_seed + + if seed is None: + seed = np.random.default_rng().integers(2**32 - 1) + # python hash seed + os.environ["PYTHONHASHSEED"] = str(seed) + # transformers seed + set_seed(seed) + + +class LocalHuggingFaceAPI(HuggingFaceAPI): + """ + Wrapper around HuggingFaceAPI that allows for quantized models to be used during + evaluation. + """ + + def __init__( # noqa: C901 + self, + model_name: str, + model: AutoModelForCausalLM, + config: GenerateConfig | None = None, + **model_args: dict[str, Any], + ) -> None: + self.model_name = model_name + self.base_url = None + self.api_key = None + self.api_key_vars = ["HF_TOKEN"] + self._apply_api_key_overrides() + + if config is None: + config = GenerateConfig() + + # set random seeds + if config.seed is not None: + set_random_seeds(config.seed) + + # collect known model_args (then delete them so we can pass the rest on) + def collect_model_arg(name: str) -> Any | None: # noqa: ANN401 + nonlocal model_args + value = model_args.get(name) + if value is not None: + model_args.pop(name) + return value + + device = collect_model_arg("device") + tokenizer = collect_model_arg("tokenizer") + model_path = collect_model_arg("model_path") + tokenizer_path = collect_model_arg("tokenizer_path") + self.batch_size = collect_model_arg("batch_size") + self.chat_template = collect_model_arg("chat_template") + self.tokenizer_call_args = collect_model_arg("tokenizer_call_args") + self.enable_thinking = collect_model_arg("enable_thinking") + if self.tokenizer_call_args is None: + self.tokenizer_call_args = {} + self.hidden_states = collect_model_arg("hidden_states") + + # device + if device: + self.device = device + elif torch.backends.mps.is_available(): + self.device = "mps" + elif torch.cuda.is_available(): + self.device = "cuda:0" + else: + self.device = "cpu" + + # model + self.model = model + + # tokenizer + if tokenizer: + self.tokenizer = AutoTokenizer.from_pretrained(tokenizer) # type: ignore[no-untyped-call] + elif model_path: + if tokenizer_path: + self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) # type: ignore[no-untyped-call] + else: + self.tokenizer = AutoTokenizer.from_pretrained(model_path) # type: ignore[no-untyped-call] + else: + self.tokenizer = AutoTokenizer.from_pretrained(model_name) # type: ignore[no-untyped-call] + # LLMs generally don't have a pad token and we need one for batching + self.tokenizer.pad_token = self.tokenizer.eos_token + self.tokenizer.padding_side = "left" + + +@modelapi(name="local_hf") +def local_hf() -> type[LocalHuggingFaceAPI]: + return LocalHuggingFaceAPI diff --git a/fouroversix/scripts/ptq/experiment.py b/fouroversix/scripts/ptq/experiment.py new file mode 100644 index 0000000..4331bff --- /dev/null +++ b/fouroversix/scripts/ptq/experiment.py @@ -0,0 +1,22 @@ +from sqlalchemy import JSON, Column, Float, Integer, String +from sqlalchemy.ext.declarative import declarative_base + +Base = declarative_base() + + +class Experiment(Base): + """A PTQ experiment with results.""" + + __tablename__ = "experiments" + + id = Column(Integer, primary_key=True, autoincrement=True) + group_name = Column(String) + model_name = Column(String, nullable=False) + task = Column(String, nullable=False) + metric_name = Column(String, nullable=False) + metric_value = Column(Float, nullable=False) + ptq_method = Column(String, nullable=False) + activation_scale_rule = Column(String, nullable=False) + weight_scale_rule = Column(String, nullable=False) + smoothquant_alpha = Column(Float, nullable=True) + results = Column(JSON, nullable=False) diff --git a/fouroversix/scripts/ptq/tasks/__init__.py b/fouroversix/scripts/ptq/tasks/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/fouroversix/scripts/ptq/tasks/wikitext_train/preprocess_wikitext.py b/fouroversix/scripts/ptq/tasks/wikitext_train/preprocess_wikitext.py new file mode 100644 index 0000000..e5dff22 --- /dev/null +++ b/fouroversix/scripts/ptq/tasks/wikitext_train/preprocess_wikitext.py @@ -0,0 +1,48 @@ +import re + + +def wikitext_detokenizer(doc): + string = doc["page"] + # contractions + string = string.replace("s '", "s'") + string = re.sub(r"/' [0-9]/", r"/'[0-9]/", string) + # number separators + string = string.replace(" @-@ ", "-") + string = string.replace(" @,@ ", ",") + string = string.replace(" @.@ ", ".") + # punctuation + string = string.replace(" : ", ": ") + string = string.replace(" ; ", "; ") + string = string.replace(" . ", ". ") + string = string.replace(" ! ", "! ") + string = string.replace(" ? ", "? ") + string = string.replace(" , ", ", ") + # double brackets + string = re.sub(r"\(\s*([^\)]*?)\s*\)", r"(\1)", string) + string = re.sub(r"\[\s*([^\]]*?)\s*\]", r"[\1]", string) + string = re.sub(r"{\s*([^}]*?)\s*}", r"{\1}", string) + string = re.sub(r"\"\s*([^\"]*?)\s*\"", r'"\1"', string) + string = re.sub(r"'\s*([^']*?)\s*'", r"'\1'", string) + # miscellaneous + string = string.replace("= = = =", "====") + string = string.replace("= = =", "===") + string = string.replace("= =", "==") + string = string.replace(" " + chr(176) + " ", chr(176)) + string = string.replace(" \n", "\n") + string = string.replace("\n ", "\n") + string = string.replace(" N ", " 1 ") + string = string.replace(" 's", "'s") + + return string + + +def process_results(doc, results): + (loglikelihood,) = results + # IMPORTANT: wikitext counts number of words in *original doc before detokenization* + _words = len(re.split(r"\s+", doc["page"])) + _bytes = len(doc["page"].encode("utf-8")) + return { + "word_perplexity": (loglikelihood, _words), + "byte_perplexity": (loglikelihood, _bytes), + "bits_per_byte": (loglikelihood, _bytes), + } diff --git a/fouroversix/scripts/ptq/tasks/wikitext_train/wikitext_train.yaml b/fouroversix/scripts/ptq/tasks/wikitext_train/wikitext_train.yaml new file mode 100644 index 0000000..ba16ab3 --- /dev/null +++ b/fouroversix/scripts/ptq/tasks/wikitext_train/wikitext_train.yaml @@ -0,0 +1,16 @@ +task: wikitext_train +dataset_path: EleutherAI/wikitext_document_level +dataset_name: wikitext-2-raw-v1 +output_type: loglikelihood_rolling +test_split: train +doc_to_text: "" +doc_to_target: !function preprocess_wikitext.wikitext_detokenizer +process_results: !function preprocess_wikitext.process_results +should_decontaminate: true +doc_to_decontamination_query: "{{page}}" +metric_list: + - metric: word_perplexity + - metric: byte_perplexity + - metric: bits_per_byte +metadata: + version: 2.0 diff --git a/fouroversix/scripts/ptq/utils.py b/fouroversix/scripts/ptq/utils.py new file mode 100644 index 0000000..5a20c1d --- /dev/null +++ b/fouroversix/scripts/ptq/utils.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +from enum import Enum + + +class EvaluationFramework(str, Enum): + """Frameworks to use when evaluating models.""" + + inspect_ai = "inspect_ai" + lm_eval = "lm_eval" + + +class PTQMethod(str, Enum): + """Methods of post-training quantization.""" + + awq = "awq" + high_precision = "high_precision" + gptq = "gptq" + rtn = "rtn" + smoothquant = "smoothquant" + spinquant = "spinquant" + + +def get_model_size(model_name: str | None) -> float: + return float(model_name.split("-")[-1][:-1]) if model_name else 0.0 diff --git a/fouroversix/scripts/resources.py b/fouroversix/scripts/resources.py new file mode 100644 index 0000000..9707616 --- /dev/null +++ b/fouroversix/scripts/resources.py @@ -0,0 +1,388 @@ +from __future__ import annotations + +import configparser +import os +import shutil +import subprocess +from enum import Enum +from pathlib import Path +from typing import TYPE_CHECKING + +import modal +import tomllib + +if TYPE_CHECKING: + from collections.abc import Callable + +FOUROVERSIX_CACHE_PATH = Path("/fouroversix") +FOUROVERSIX_INSTALL_PATH = Path("/root/fouroversix") +KERNEL_DEV_MODE = os.getenv("KERNEL_DEV_MODE", "0") == "1" + +app = modal.App("fouroversix") +cache_volume = modal.Volume.from_name("fouroversix", create_if_missing=True) +hf_secret = modal.Secret.from_name("huggingface-secret") +wandb_secret = modal.Secret.from_name("wandb-secret") + + +class Dependency(str, Enum): + """Dependencies to add to the base image.""" + + awq = "awq" + fast_hadamard_transform = "fast_hadamard_transform" + flame = "flame" + flash_attention = "flash_attention" + fouroversix = "fouroversix" + fp_quant = "fp_quant" + qutlass = "qutlass" + spinquant = "spinquant" + transformer_engine = "transformer_engine" + + +class Submodule(str, Enum): + """Submodules of Four Over Six to add to the base image.""" + + cutlass = "cutlass" + flame = "flame" + fast_hadamard_transform = "fast_hadamard_transform" + fp_quant = "fp_quant" + llm_awq = "llm_awq" + qutlass = "qutlass" + spinquant = "spinquant" + + def has_untracked_or_unstaged_changes(self) -> bool: + """Check if the submodule has untracked or unstaged changes.""" + + git_status = subprocess.run( + [ + "git", + "-C", + self.get_local_path(), + "status", + "--porcelain", + ], + check=False, + stderr=subprocess.DEVNULL, + stdout=subprocess.PIPE, + text=True, + ) + + return bool(git_status.stdout.strip()) + + def get_install_path(self) -> str: + """Get the path where this submodule will be installed in the Modal image.""" + return f"{FOUROVERSIX_INSTALL_PATH}/{self.get_local_path()}" + + def get_local_path(self) -> str: + """Get the path of the submodule relative to the root directory.""" + return f"third_party/{self.value.replace('_', '-')}" + + def get_remote_url(self) -> str: + """Get the remote URL of the submodule.""" + + gitmodules_path = Path(__file__).parent.parent / ".gitmodules" + + if not gitmodules_path.exists(): + gitmodules_path = FOUROVERSIX_INSTALL_PATH / ".gitmodules" + + with gitmodules_path.open() as f: + # Remove leading whitespace to make it a valid INI file + gitmodules_contents = "\n".join(line.lstrip() for line in f.readlines()) + + config = configparser.ConfigParser() + config.read_string(gitmodules_contents) + + for section in config.sections(): + if config[section]["path"] == self.get_local_path(): + url = config[section]["url"] + break + + if url.startswith("https://"): + return url + + msg = f"Unsupported remote URL format: {url}" + raise ValueError(msg) + + +cuda_version_to_image_tag = { + "12.8": "nvcr.io/nvidia/cuda-dl-base:25.03-cuda12.8-devel-ubuntu24.04", + "12.9": "nvcr.io/nvidia/cuda-dl-base:25.06-cuda12.9-devel-ubuntu24.04", + "13.0": "nvcr.io/nvidia/cuda-dl-base:25.09-cuda13.0-devel-ubuntu24.04", + "13.1": "nvcr.io/nvidia/cuda-dl-base:25.12-cuda13.1-devel-ubuntu24.04", +} + + +def add_submodule(img: modal.Image, submodule: Submodule) -> modal.Image: + if submodule.has_untracked_or_unstaged_changes(): + # Submodule has uncommitted changes, build image with local copy + return img.add_local_dir( + submodule.get_local_path(), + submodule.get_install_path(), + copy=True, + ) + + # Submodule has no uncommitted changes, download from remote to save time + return img.run_commands( + f"git clone {submodule.get_remote_url()} {submodule.get_install_path()}", + ) + + +def install_flash_attn() -> None: + subprocess.run( + ["pip", "install", "flash-attn", "--no-build-isolation"], + check=False, + ) + + +def install_fouroversix() -> None: + subprocess.run( + [ + "pip", + "install", + "--no-deps", + "--no-build-isolation", + "-e", + FOUROVERSIX_INSTALL_PATH.as_posix(), + ], + check=False, + ) + + +def install_fouroversix_non_editable() -> None: + shutil.copytree( + FOUROVERSIX_CACHE_PATH / "build", + FOUROVERSIX_INSTALL_PATH / "build", + ) + subprocess.run( + ["python", "setup.py", "build_ext", "--inplace"], + check=False, + ) + shutil.copytree( + FOUROVERSIX_INSTALL_PATH / "build", + FOUROVERSIX_CACHE_PATH / "build", + dirs_exist_ok=True, + ) + subprocess.run( + [ + "pip", + "install", + "--no-deps", + "--no-build-isolation", + FOUROVERSIX_INSTALL_PATH.as_posix(), + ], + check=False, + ) + + +def install_qutlass() -> None: + subprocess.run( + [ + "pip", + "install", + "--no-build-isolation", + Submodule.qutlass.get_install_path(), + ], + check=False, + ) + + +def get_image( # noqa: C901, PLR0912 + dependencies: list[Dependency] | None = None, + *, + cuda_version: str = "12.9", + deploy: bool = False, + extra_env: dict[str, str] | None = None, + extra_pip_dependencies: list[str] | None = None, + include_tests: bool = False, + python_version: str = "3.13", + pytorch_version: str = "2.10.0", + run_before_copy: Callable[[modal.Image], modal.Image] | None = None, +) -> modal.Image: + if dependencies is None: + dependencies = [Dependency.fouroversix] + + pyproject_path = Path(__file__).parent.parent / "pyproject.toml" + + if not pyproject_path.exists(): + pyproject_path = Path(__file__).parent.parent / "fouroversix" / "pyproject.toml" + + with pyproject_path.open("rb") as f: + pyproject_data = tomllib.load(f) + + img = ( + modal.Image.from_registry( + cuda_version_to_image_tag[cuda_version], + add_python=python_version, + ) + .entrypoint([]) + .apt_install("clang", "git") + .uv_pip_install(*pyproject_data["build-system"]["requires"], "numpy") + .uv_pip_install( + f"torch=={pytorch_version}", + extra_index_url=( + f"https://download.pytorch.org/whl/cu{cuda_version.replace('.', '')}" + ), + ) + ) + + for dependency in dependencies: + if dependency == Dependency.awq: + img = add_submodule(img, Submodule.llm_awq).run_commands( + f"pip install --no-deps -e {Submodule.llm_awq.get_install_path()}", + ) + + if dependency == Dependency.fast_hadamard_transform: + img = add_submodule(img, Submodule.fast_hadamard_transform).run_commands( + f"pip install {Submodule.fast_hadamard_transform.get_install_path()} " + "--no-build-isolation", + ) + + if dependency == Dependency.flame: + img = ( + img.apt_install("pciutils") + .uv_pip_install( + "flash-linear-attention", + "ninja", + "psutil", + "git+https://github.com/pytorch/torchtitan.git@0b44d4c", + "tyro", + "wheel", + ) + .run_commands( + "git clone https://github.com/fla-org/flame.git " + f"{FOUROVERSIX_INSTALL_PATH}/third_party/flame", + f"pip install -e {FOUROVERSIX_INSTALL_PATH}/third_party/flame", + ) + ) + + if dependency == Dependency.flash_attention: + img = img.run_function( + install_flash_attn, + cpu=64, + memory=128 * 1024, + gpu="B200", + ) + + if dependency == Dependency.fouroversix: + img = ( + add_submodule( + img.env( + {"CUDA_ARCHS": "100", "FORCE_BUILD": "1", "MAX_JOBS": "32"}, + ), + Submodule.cutlass, + ) + .add_local_file( + "pyproject.toml", + f"{FOUROVERSIX_INSTALL_PATH}/pyproject.toml", + copy=True, + ) + .uv_pip_install( + *pyproject_data["project"]["optional-dependencies"]["evals"], + ) + .add_local_file( + "setup.py", + f"{FOUROVERSIX_INSTALL_PATH}/setup.py", + copy=True, + ) + .add_local_file( + "src/fouroversix/__init__.py", + f"{FOUROVERSIX_INSTALL_PATH}/src/fouroversix/__init__.py", + copy=True, + ) + ) + + if KERNEL_DEV_MODE: + img = ( + img.add_local_file( + "README.md", + f"{FOUROVERSIX_INSTALL_PATH}/README.md", + copy=True, + ) + .add_local_file( + "LICENSE.md", + f"{FOUROVERSIX_INSTALL_PATH}/LICENSE.md", + copy=True, + ) + .workdir(FOUROVERSIX_INSTALL_PATH) + ) + + img = img.add_local_dir( + "src/fouroversix/csrc", + f"{FOUROVERSIX_INSTALL_PATH}/src/fouroversix/csrc", + copy=True, + ) + + if not KERNEL_DEV_MODE: + img = img.run_function(install_fouroversix, cpu=32, memory=64 * 1024) + + if dependency == Dependency.fp_quant: + img = add_submodule(img, Submodule.fp_quant).run_commands( + f"pip install {Submodule.fp_quant.get_install_path()}/inference_lib", + ) + + if dependency == Dependency.qutlass: + img = ( + add_submodule(img.apt_install("cmake"), Submodule.qutlass) + .env({"MAX_JOBS": "32"}) + .run_function(install_qutlass, gpu="B200", cpu=32, memory=64 * 1024) + ) + + if dependency == Dependency.spinquant: + img = add_submodule(img, Submodule.spinquant) + + if dependency == Dependency.transformer_engine: + img = img.uv_pip_install( + "transformer_engine[pytorch]", + extra_options="--no-build-isolation", + ) + + if extra_pip_dependencies is not None: + img = img.uv_pip_install(*extra_pip_dependencies) + + img = img.env({"HF_HOME": FOUROVERSIX_CACHE_PATH.as_posix(), **(extra_env or {})}) + + if run_before_copy is not None: + img = run_before_copy(img) + + # Add source files after all dependencies are added so we can avoid rebuilding when + # they change + for dependency in dependencies: + if dependency == Dependency.flame: + img = ( + img.add_local_dir( + "third_party/flame/custom_models", + f"{FOUROVERSIX_INSTALL_PATH}/third_party/flame/custom_models", + ) + .add_local_dir( + "scripts/train/configs", + f"{FOUROVERSIX_INSTALL_PATH}/scripts/train/configs", + ) + .add_local_file( + "third_party/flame/train.sh", + f"{FOUROVERSIX_INSTALL_PATH}/third_party/flame/train.sh", + ) + ) + + if dependency == Dependency.fouroversix: + img = img.add_local_dir( + "src", + f"{FOUROVERSIX_INSTALL_PATH}/src", + copy=deploy or KERNEL_DEV_MODE, + ignore=lambda p: p.suffix == ".so", + ).add_local_file( + ".gitmodules", + f"{FOUROVERSIX_INSTALL_PATH}/.gitmodules", + copy=deploy or KERNEL_DEV_MODE, + ) + + if KERNEL_DEV_MODE: + img = img.run_function( + install_fouroversix_non_editable, + cpu=32, + memory=64 * 1024, + volumes={FOUROVERSIX_CACHE_PATH.as_posix(): cache_volume}, + ) + + if include_tests: + img = img.add_local_dir("tests", f"{FOUROVERSIX_INSTALL_PATH}/tests") + + return img diff --git a/fouroversix/scripts/speedtest/__init__.py b/fouroversix/scripts/speedtest/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/fouroversix/scripts/speedtest/matmul.py b/fouroversix/scripts/speedtest/matmul.py new file mode 100644 index 0000000..64edd52 --- /dev/null +++ b/fouroversix/scripts/speedtest/matmul.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +from typing import Any + +import click +import modal + +from ..resources import app, get_image + +img = get_image() + +with img.imports(): + import torch + import torch.utils.benchmark as benchmark + from fouroversix import DataType, MatmulBackend, QuantizationConfig, quantize_to_fp4 + from fouroversix.matmul.frontend import AVAILABLE_BACKENDS + + +def run_speedtest( + *, + dtype: DataType = DataType.nvfp4, + m: int = 1024, + n: int = 1024, + k: int = 1024, + repeats: int = 100, +) -> None: + """Test speed on a B200 on Modal.""" + + x = torch.randn(m, k, dtype=torch.bfloat16, device="cuda") + y = torch.randn(k, n, dtype=torch.bfloat16, device="cuda") + + config = QuantizationConfig(dtype=dtype) + x_quantized = quantize_to_fp4(x, config) + y_quantized = quantize_to_fp4(y, config) + out_dtype = DataType.bfloat16 + + print(f"Testing with {m}x{k} @ {k}x{n}") + + for backend in [MatmulBackend.cutlass, MatmulBackend.pytorch]: + backend_cls = AVAILABLE_BACKENDS[backend] + print(f"{backend.value}: ", end="") + + if not backend_cls.is_available(): + print("Not available") + continue + + if not backend_cls.is_supported( + x_quantized, + y_quantized, + out_dtype=out_dtype, + ): + print("Not supported") + continue + + t = benchmark.Timer( + setup="from fouroversix import fp4_matmul", + stmt=( + "fp4_matmul(x_quantized, y_quantized, backend=backend, " + "out_dtype=out_dtype)" + ), + globals={ + "x_quantized": x_quantized, + "y_quantized": y_quantized, + "backend": backend, + "out_dtype": out_dtype, + }, + ) + + print(f"{t.timeit(repeats).mean * 1000:.4f}ms") + + +@app.function(image=img, cpu=4, memory=8 * 1024, gpu="B200") +def run_speedtest_on_modal(**kwargs: dict[str, Any]) -> None: + run_speedtest(**kwargs) + + +@click.command() +@click.option("--dtype", type=DataType, default=DataType.nvfp4) +@click.option("--m", type=int, default=1024) +@click.option("--modal", is_flag=True) +@click.option("--n", type=int, default=1024) +@click.option("--k", type=int, default=1024) +@click.option("--repeats", type=int, default=100) +def cli(**kwargs: dict[str, Any]) -> None: + if kwargs.pop("modal"): + with modal.enable_output(), app.run(): + run_speedtest_on_modal.remote(**kwargs) + else: + run_speedtest(**kwargs) + + +if __name__ == "__main__": + cli() diff --git a/fouroversix/scripts/speedtest/quantize.py b/fouroversix/scripts/speedtest/quantize.py new file mode 100644 index 0000000..5192376 --- /dev/null +++ b/fouroversix/scripts/speedtest/quantize.py @@ -0,0 +1,108 @@ +from __future__ import annotations + +from typing import Any + +import click +import modal + +from ..resources import Dependency, app, get_image + +img = get_image(dependencies=[Dependency.transformer_engine, Dependency.fouroversix]) + +with img.imports(): + import torch + import torch.utils.benchmark as benchmark + from fouroversix import QuantizationConfig, QuantizeBackend, RoundStyle, ScaleRule + from fouroversix.quantize.frontend import AVAILABLE_BACKENDS + + +def run_speedtest( + *, + block_scale_2d: bool = False, + input_shape: str = "1024,1024", + repeats: int = 100, + rht: bool = False, + round_style: str = "nearest", + scale_rule: str = "mse", + transpose: bool = False, +) -> None: + """Test speed on a B200 on Modal.""" + + input_shape = tuple(int(dim.strip()) for dim in input_shape.split(",")) + x = torch.randn(input_shape, dtype=torch.bfloat16, device="cuda") + + print("Testing with config:") + print(f"- block_scale_2d: {block_scale_2d}") + print(f"- input_shape: {input_shape}") + print(f"- rht: {rht}") + print(f"- round_style: {round_style}") + print(f"- scale_rule: {scale_rule}") + print(f"- transpose: {transpose}") + print() + + for backend in [ + QuantizeBackend.cuda, + QuantizeBackend.transformer_engine, + QuantizeBackend.triton, + QuantizeBackend.pytorch, + ]: + config = QuantizationConfig( + backend=backend, + block_scale_2d=block_scale_2d, + rht=rht, + round_style=RoundStyle(round_style), + scale_rule=ScaleRule(scale_rule), + transpose=transpose, + ) + + backend_cls = AVAILABLE_BACKENDS[backend] + print(f"{backend.value}: ", end="") + + if not backend_cls.is_available(): + print("Not available") + continue + + if not backend_cls.is_supported(x, config): + print("Not supported") + continue + + config = QuantizationConfig( + backend=backend, + rht=rht, + round_style=RoundStyle(round_style), + scale_rule=ScaleRule(scale_rule), + ) + + t = benchmark.Timer( + setup="from fouroversix import quantize_to_fp4", + stmt="quantize_to_fp4(x, config)", + globals={"x": x, "config": config}, + ) + + print(f"{t.timeit(repeats).mean * 1000:.4f}ms") + + +@app.function(image=img, cpu=4, memory=8 * 1024, gpu="B200") +def run_speedtest_on_modal(**kwargs: dict[str, Any]) -> None: + run_speedtest(**kwargs) + + +@click.command() +@click.option("--block-scale-2d", is_flag=True) +@click.option("--input-shape", type=str, default="1024,1024") +@click.option("--modal", is_flag=True) +@click.option("--repeats", type=int, default=100) +@click.option("--rht", is_flag=True) +@click.option("--round-style", type=RoundStyle, default=RoundStyle.nearest) +@click.option("--scale-rule", type=ScaleRule, default=ScaleRule.mse) +@click.option("--transpose", is_flag=True) +def cli(**kwargs: dict[str, Any]) -> None: + if kwargs.pop("modal"): + with modal.enable_output(), app.run(): + run_speedtest_on_modal.remote(**kwargs) + else: + run_speedtest(**kwargs) + + +if __name__ == "__main__": + cli() diff --git a/fouroversix/scripts/test_on_modal.py b/fouroversix/scripts/test_on_modal.py new file mode 100644 index 0000000..eb29804 --- /dev/null +++ b/fouroversix/scripts/test_on_modal.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +from pathlib import Path + +from .resources import FOUROVERSIX_INSTALL_PATH, Dependency, app, get_image + +img = get_image( + dependencies=[Dependency.transformer_engine, Dependency.fouroversix], + include_tests=True, +) + +with img.imports(): + import pytest + + +@app.function(image=img, cpu=4, memory=8 * 1024, gpu="B200", timeout=30 * 60) +def run_tests(*args: list[str]) -> None: + """Run tests on a B200 on Modal.""" + + args = list(args) + tests_path = (Path(FOUROVERSIX_INSTALL_PATH) / "tests").as_posix() + + if len(args) == 0: + args = [tests_path] + elif "tests" in args[0] or "test_" in args[0]: + args[0] = (Path(FOUROVERSIX_INSTALL_PATH) / args[0]).as_posix() + else: + args = [tests_path, *args] + + pytest.main(args) diff --git a/fouroversix/scripts/train/__init__.py b/fouroversix/scripts/train/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/fouroversix/scripts/train/__main__.py b/fouroversix/scripts/train/__main__.py new file mode 100644 index 0000000..60ab80d --- /dev/null +++ b/fouroversix/scripts/train/__main__.py @@ -0,0 +1,244 @@ +from __future__ import annotations + +import os +import subprocess +import time +from pathlib import Path +from typing import Any + +import click +import modal + +from ..resources import ( + FOUROVERSIX_CACHE_PATH, + FOUROVERSIX_INSTALL_PATH, + Dependency, + app, + cache_volume, + get_image, + wandb_secret, +) + +img = get_image( + dependencies=[Dependency.flash_attention, Dependency.flame, Dependency.fouroversix], + extra_pip_dependencies=["datasets<4.6"], +) + + +def train( + *, + batch_size: int, + checkpoint_interval: int, + checkpoint_keep_latest_k: int, + checkpoint_load_step: int, + context_length: int, + dataset: str, + dataset_name: str, + dataset_split: str, + exp_folder: str, + gradient_accumulation_steps: int, + initial_load_path: str | None, + job_config_file: str, + lr: float, + lr_decay_type: str, + model_config: str, + model_name: str, + no_torch_compile: bool, + seed: int, + tokenizer: str, + training_steps: int | None, +) -> None: + import torch + + # Cache activations and gradients and set dump folder + os.environ["CACHE_ACTIVATIONS"] = "1" + os.environ["CACHE_GRADIENTS"] = "1" + os.environ["DUMP_FOLDER"] = f"{exp_folder}/{model_name}" + + # Set MODEL_NAME for wandb + os.environ["MODEL_NAME"] = model_name + + # Set NGPU for flame + os.environ["NGPU"] = str(torch.cuda.device_count()) + + if training_steps is None: + if dataset_name == "sample-10BT": + num_tokens = 10_000_000_000 + elif dataset_name == "sample-100BT": + num_tokens = 100_000_000_000 + else: + msg = ( + "You must provide the number of training steps if not using the " + "sample-10BT or sample-100BT datasets" + ) + raise ValueError(msg) + + training_steps = num_tokens // int( + context_length * batch_size * torch.cuda.device_count(), + ) + + # Start training + args = [ + "bash", + "train.sh", + "--job.config_file", + job_config_file, + "--job.dump_folder", + f"{exp_folder}/{model_name}", + "--model.config", + model_config, + "--model.tokenizer_path", + tokenizer, + "--optimizer.name", + "AdamW", + "--optimizer.lr", + str(lr), + "--lr_scheduler.warmup_steps", + "0", + "--lr_scheduler.decay_ratio", + "0.15", + "--lr_scheduler.decay_type", + lr_decay_type, + "--lr_scheduler.lr_min", + "0.01", + "--training.batch_size", + "1", + "--training.seq_len", + str(int(context_length * batch_size)), + "--training.context_len", + str(context_length), + "--training.varlen", + "--training.gradient_accumulation_steps", + str(gradient_accumulation_steps), + "--training.steps", + str(training_steps), + "--training.max_norm", + "1.0", + "--training.skip_nan_inf", + "--training.dataset", + dataset, + "--training.dataset_name", + dataset_name, + "--training.dataset_split", + dataset_split, + "--training.num_workers", + "32", + "--training.prefetch_factor", + "2", + "--training.seed", + str(seed), + "--checkpoint.interval", + str(checkpoint_interval), + "--checkpoint.load_step", + str(checkpoint_load_step), + "--checkpoint.keep_latest_k", + str(checkpoint_keep_latest_k), + "--metrics.log_freq", + "1", + ] + + if not no_torch_compile: + args.append("--training.compile") + + if initial_load_path is not None: + args.extend( + [ + "--checkpoint.initial_load_path", + initial_load_path, + "--checkpoint.no_initial_load_model_weights_only", + ], + ) + + subprocess.run(args, check=True) + + +@app.cls( + image=img, + gpu="B200:8", + timeout=24 * 60 * 60, + cpu=64, + memory=8 * 64 * 1024, + volumes={FOUROVERSIX_CACHE_PATH: cache_volume}, + secrets=[wandb_secret], +) +class ModalTrainer: + """Run training jobs on Modal.""" + + @modal.method() + def train(self, **kwargs: dict[str, Any]) -> None: + """Start a training job on Modal.""" + os.chdir(FOUROVERSIX_INSTALL_PATH / "third_party" / "flame") + train(**kwargs) + + +@click.command() +@click.option("--batch-size", type=float, default=16) +@click.option("--checkpoint-interval", type=int, default=1000) +@click.option("--checkpoint-keep-latest-k", type=int, default=0) +@click.option("--checkpoint-load-step", type=int, default=-1) +@click.option("--context-length", type=int, default=8192) +@click.option("--dataset", type=str, default="HuggingFaceFW/fineweb-edu") +@click.option("--dataset-name", type=str, default="sample-100BT") +@click.option("--dataset-split", type=str, default="train") +@click.option("--detach", is_flag=True) +@click.option("--exp-folder", type=str, default="exp") +@click.option("--gradient-accumulation-steps", type=int, default=1) +@click.option("--initial-load-path", type=str) +@click.option("--job-config-file", type=str, default="flame/models/fla.toml") +@click.option("--lr", type=float, default=1.2e-3) +@click.option("--lr-decay-type", type=str, default="linear") +@click.option("--modal", is_flag=True) +@click.option("--modal-gpu", type=str, default="B200:8") +@click.option("--model-config", type=str, required=True) +@click.option("--model-name", type=str, required=True) +@click.option("--no-torch-compile", is_flag=True) +@click.option("--seed", type=int, default=42) +@click.option("--tokenizer", type=str, default="fla-hub/transformer-1.3B-100B") +@click.option("--training-steps", type=int, default=None) +@click.option("--wait-for-pid", type=int, default=None) +def cli(**kwargs: dict[str, Any]) -> None: + # Options that are not passed to the train function + detach = kwargs.pop("detach", False) + modal_gpu = kwargs.pop("modal_gpu", "B200:8") + use_modal = kwargs.pop("modal", False) + wait_for_pid = kwargs.pop("wait_for_pid", None) + + # Wait for the previous training job to finish + if wait_for_pid is not None: + while ( + subprocess.run(["kill", "-0", str(wait_for_pid)], check=False).returncode + == 0 + ): + time.sleep(1) + + time.sleep(60) + + if not Path(kwargs["model_config"]).exists(): + kwargs["model_config"] = ( + Path(__file__).parent.parent.parent / kwargs["model_config"] + ) + + if not Path(kwargs["model_config"]).exists(): + msg = f"Model config file not found: {kwargs['model_config']}" + raise FileNotFoundError(msg) + + # Set exp folder on Modal + if use_modal: + with modal.enable_output(), app.run(detach=detach): + kwargs["exp_folder"] = (FOUROVERSIX_CACHE_PATH / "exp").as_posix() + kwargs["model_config"] = ( + (FOUROVERSIX_INSTALL_PATH / kwargs["model_config"]) + .absolute() + .as_posix() + ) + + ModalTrainer.with_options(gpu=modal_gpu)().train.remote(**kwargs) + else: + kwargs["model_config"] = Path(kwargs["model_config"]).absolute().as_posix() + + os.chdir(Path(__file__).parent.parent.parent / "third_party" / "flame") + train(**kwargs) + + +if __name__ == "__main__": + cli() diff --git a/fouroversix/scripts/train/configs/transformer_1B_bf16.json b/fouroversix/scripts/train/configs/transformer_1B_bf16.json new file mode 100644 index 0000000..448abc7 --- /dev/null +++ b/fouroversix/scripts/train/configs/transformer_1B_bf16.json @@ -0,0 +1,26 @@ +{ + "attention_bias": false, + "bos_token_id": 1, + "eos_token_id": 2, + "fuse_cross_entropy": true, + "fuse_norm": true, + "fuse_swiglu": false, + "hidden_act": "swish", + "hidden_ratio": 4, + "hidden_size": 2048, + "initializer_range": 0.02, + "max_position_embeddings": 8192, + "model_type": "fp4_transformer", + "num_heads": 32, + "num_hidden_layers": 24, + "norm_eps": 1e-06, + "tie_word_embeddings": false, + "use_cache": true, + "vocab_size": 32000, + "qk_norm": true, + "layer_precision_configs": [ + { + "repeats": 24 + } + ] +} \ No newline at end of file diff --git a/fouroversix/scripts/train/configs/transformer_1B_fp4.json b/fouroversix/scripts/train/configs/transformer_1B_fp4.json new file mode 100644 index 0000000..0e6dc86 --- /dev/null +++ b/fouroversix/scripts/train/configs/transformer_1B_fp4.json @@ -0,0 +1,31 @@ +{ + "attention_bias": false, + "bos_token_id": 1, + "eos_token_id": 2, + "fuse_cross_entropy": true, + "fuse_norm": true, + "fuse_swiglu": false, + "hidden_act": "swish", + "hidden_ratio": 4, + "hidden_size": 2048, + "initializer_range": 0.02, + "max_position_embeddings": 8192, + "model_type": "fp4_transformer", + "num_heads": 32, + "num_hidden_layers": 24, + "norm_eps": 1e-06, + "tie_word_embeddings": false, + "use_cache": true, + "vocab_size": 32000, + "qk_norm": true, + "layer_precision_configs": [ + { + "repeats": 20, + "dtype": "nvfp4", + "scale_rule": "static_6" + }, + { + "repeats": 4 + } + ] +} \ No newline at end of file diff --git a/fouroversix/scripts/train/configs/transformer_1B_fp4_fouroversix.json b/fouroversix/scripts/train/configs/transformer_1B_fp4_fouroversix.json new file mode 100644 index 0000000..d0cccc4 --- /dev/null +++ b/fouroversix/scripts/train/configs/transformer_1B_fp4_fouroversix.json @@ -0,0 +1,31 @@ +{ + "attention_bias": false, + "bos_token_id": 1, + "eos_token_id": 2, + "fuse_cross_entropy": true, + "fuse_norm": true, + "fuse_swiglu": false, + "hidden_act": "swish", + "hidden_ratio": 4, + "hidden_size": 2048, + "initializer_range": 0.02, + "max_position_embeddings": 8192, + "model_type": "fp4_transformer", + "num_heads": 32, + "num_hidden_layers": 24, + "norm_eps": 1e-06, + "tie_word_embeddings": false, + "use_cache": true, + "vocab_size": 32000, + "qk_norm": true, + "layer_precision_configs": [ + { + "repeats": 20, + "dtype": "nvfp4", + "scale_rule": "mse" + }, + { + "repeats": 4 + } + ] +} \ No newline at end of file diff --git a/fouroversix/scripts/train/configs/transformer_340M_bf16.json b/fouroversix/scripts/train/configs/transformer_340M_bf16.json new file mode 100644 index 0000000..1df0d5a --- /dev/null +++ b/fouroversix/scripts/train/configs/transformer_340M_bf16.json @@ -0,0 +1,25 @@ +{ + "attention_bias": false, + "bos_token_id": 1, + "eos_token_id": 2, + "fuse_cross_entropy": true, + "fuse_norm": true, + "fuse_swiglu": false, + "hidden_act": "swish", + "hidden_size": 1024, + "initializer_range": 0.02, + "max_position_embeddings": 8192, + "model_type": "fp4_transformer", + "num_heads": 16, + "num_hidden_layers": 24, + "norm_eps": 1e-06, + "tie_word_embeddings": false, + "use_cache": true, + "vocab_size": 32000, + "qk_norm": true, + "layer_precision_configs": [ + { + "repeats": 24 + } + ] +} \ No newline at end of file diff --git a/fouroversix/scripts/train/configs/transformer_340M_fp4.json b/fouroversix/scripts/train/configs/transformer_340M_fp4.json new file mode 100644 index 0000000..597b6d4 --- /dev/null +++ b/fouroversix/scripts/train/configs/transformer_340M_fp4.json @@ -0,0 +1,30 @@ +{ + "attention_bias": false, + "bos_token_id": 1, + "eos_token_id": 2, + "fuse_cross_entropy": true, + "fuse_norm": true, + "fuse_swiglu": false, + "hidden_act": "swish", + "hidden_size": 1024, + "initializer_range": 0.02, + "max_position_embeddings": 8192, + "model_type": "fp4_transformer", + "num_heads": 16, + "num_hidden_layers": 24, + "norm_eps": 1e-06, + "tie_word_embeddings": false, + "use_cache": true, + "vocab_size": 32000, + "qk_norm": true, + "layer_precision_configs": [ + { + "repeats": 20, + "dtype": "nvfp4", + "scale_rule": "static_6" + }, + { + "repeats": 4 + } + ] +} \ No newline at end of file diff --git a/fouroversix/scripts/train/configs/transformer_340M_fp4_fouroversix.json b/fouroversix/scripts/train/configs/transformer_340M_fp4_fouroversix.json new file mode 100644 index 0000000..3ae5463 --- /dev/null +++ b/fouroversix/scripts/train/configs/transformer_340M_fp4_fouroversix.json @@ -0,0 +1,30 @@ +{ + "attention_bias": false, + "bos_token_id": 1, + "eos_token_id": 2, + "fuse_cross_entropy": true, + "fuse_norm": true, + "fuse_swiglu": false, + "hidden_act": "swish", + "hidden_size": 1024, + "initializer_range": 0.02, + "max_position_embeddings": 8192, + "model_type": "fp4_transformer", + "num_heads": 16, + "num_hidden_layers": 24, + "norm_eps": 1e-06, + "tie_word_embeddings": false, + "use_cache": true, + "vocab_size": 32000, + "qk_norm": true, + "layer_precision_configs": [ + { + "repeats": 20, + "dtype": "nvfp4", + "scale_rule": "mse" + }, + { + "repeats": 4 + } + ] +} \ No newline at end of file diff --git a/fouroversix/scripts/train/prepare_dataset.py b/fouroversix/scripts/train/prepare_dataset.py new file mode 100644 index 0000000..f1cb9d9 --- /dev/null +++ b/fouroversix/scripts/train/prepare_dataset.py @@ -0,0 +1,15 @@ +from ..resources import FOUROVERSIX_CACHE_PATH, app, cache_volume, get_image + +img = get_image(dependencies=[], extra_pip_dependencies=["datasets"]) + +with img.imports(): + from datasets import load_dataset + + +@app.function( + image=img, + timeout=24 * 60 * 60, + volumes={FOUROVERSIX_CACHE_PATH: cache_volume}, +) +def prepare_dataset(path: str, name: str) -> None: + load_dataset(path, name) diff --git a/fouroversix/setup.py b/fouroversix/setup.py new file mode 100644 index 0000000..d761d39 --- /dev/null +++ b/fouroversix/setup.py @@ -0,0 +1,296 @@ +import functools +import os +import platform +import subprocess +import sys +import urllib.request +import warnings +from pathlib import Path +from typing import Any + +import torch +from packaging.version import Version, parse +from setuptools import setup +from setuptools.command.bdist_wheel import bdist_wheel +from torch.utils.cpp_extension import CUDA_HOME, BuildExtension, CUDAExtension + +BASE_WHEEL_URL = "https://github.com/mit-han-lab/fouroversix/releases/download" +PACKAGE_NAME = "fouroversix" +PACKAGE_VERSION = "1.1.0" + +CUTLASS_DEBUG = os.getenv("CUTLASS_DEBUG", "0") == "1" +FORCE_BUILD = os.getenv("FORCE_BUILD", "0") == "1" +FORCE_CXX11_ABI = os.getenv("FORCE_CXX11_ABI", "0") == "1" +SKIP_CUDA_BUILD = os.getenv("SKIP_CUDA_BUILD", "0") == "1" + + +@functools.cache +def get_cuda_archs() -> list[str]: + return os.getenv("CUDA_ARCHS", "100;103;110;120").split(";") + + +def get_cuda_bare_metal_version() -> Version | None: + if CUDA_HOME is None: + warnings.warn( + "nvcc was not found. Are you sure your environment has nvcc available? If " + "you're installing within a container from " + "https://hub.docker.com/r/pytorch/pytorch, only images with 'devel' in " + "their name will provide nvcc.", + stacklevel=1, + ) + return None + + raw_output = subprocess.check_output( + [CUDA_HOME + "/bin/nvcc", "-V"], + universal_newlines=True, + ) + output = raw_output.split() + release_idx = output.index("release") + 1 + return parse(output[release_idx].split(",")[0]) + + +def get_cuda_gencodes() -> list[str]: + """ + Add -gencode flags based on nvcc capabilities. + + Uses the following rules: + - sm_100/120 on CUDA >= 12.8 + - Use 100f on CUDA >= 12.9 (Blackwell family-specific) + - Map requested 110 -> 101 if CUDA < 13.0 (Thor rename) + - Embed PTX for newest arch for forward compatibility + """ + + archs = set(get_cuda_archs()) + cuda_version = get_cuda_bare_metal_version() + cc_flags = [] + + # Blackwell requires >= 12.8 + if cuda_version is not None and cuda_version >= Version("12.8"): + if "100" in archs: + cc_flags += ["-gencode", "arch=compute_100a,code=sm_100a"] + + if "103" in archs: + cc_flags += ["-gencode", "arch=compute_103a,code=sm_103a"] + + # Thor rename: 12.9 uses sm_101; 13.0+ uses sm_110 + if "110" in archs: + if cuda_version >= Version("13.0"): + cc_flags += ["-gencode", "arch=compute_110f,code=sm_110"] + elif cuda_version >= Version("12.9"): + # Provide Thor support for CUDA 12.9 via sm_101 + cc_flags += ["-gencode", "arch=compute_101f,code=sm_101"] + # else: no Thor support in older toolkits + + if "120" in archs: + # sm_120 is supported in CUDA 12.8/12.9+ toolkits + if cuda_version >= Version("12.9"): + cc_flags += ["-gencode", "arch=compute_120f,code=sm_120"] + else: + cc_flags += ["-gencode", "arch=compute_120a,code=sm_120a"] + + return cc_flags + + +def get_platform() -> str: + if sys.platform.startswith("linux"): + return f"linux_{platform.uname().machine}" + if sys.platform == "darwin": + mac_version = ".".join(platform.mac_ver()[0].split(".")[:2]) + return f"macosx_{mac_version}_x86_64" + if sys.platform == "win32": + return "win_amd64" + + msg = f"Unsupported platform: {sys.platform}" + raise ValueError(msg) + + +def get_wheel_url() -> tuple[str, str]: + torch_version_raw = parse(torch.__version__) + python_version = f"cp{sys.version_info.major}{sys.version_info.minor}" + platform_name = get_platform() + torch_version = f"{torch_version_raw.major}.{torch_version_raw.minor}" + cxx11_abi = str(torch._C._GLIBCXX_USE_CXX11_ABI).upper() # noqa: SLF001 + + # We only compile for CUDA 12.8 to save CI time. Minor versions should be + # compatible. + torch_cuda_version = parse("12.8") + cuda_version = f"cu{torch_cuda_version.major}" + + wheel_filename = ( + f"{PACKAGE_NAME}-{PACKAGE_VERSION}+{cuda_version}torch{torch_version}" + f"cxx11abi{cxx11_abi}-{python_version}-{python_version}-{platform_name}.whl" + ) + + return f"{BASE_WHEEL_URL}/v{PACKAGE_VERSION}/{wheel_filename}", wheel_filename + + +class CachedWheelsCommand(bdist_wheel): + """ + Custom bdist wheel command that checks for pre-built wheels on GitHub Releases. + + The CachedWheelsCommand plugs into the default bdist wheel, which is ran by pip + when it cannot find an existing wheel (which is currently the case for all + fouroversix installs). We use the environment parameters to detect whether there is + already a pre-built version of a compatible wheel available and short-circuits the + standard full build pipeline. + + Credit: https://github.com/Dao-AILab/flash-attention/blob/main/setup.py + """ + + def run(self) -> None: + """Run the command.""" + + if FORCE_BUILD: + return super().run() + + wheel_url, wheel_filename = get_wheel_url() + print(f"Guessing wheel URL: {wheel_url}") + + try: + urllib.request.urlretrieve(wheel_url, wheel_filename) # noqa: S310 + + # Make the archive + # Lifted from the root wheel processing command + # https://github.com/pypa/wheel/blob/cf71108ff9f6ffc36978069acb28824b44ae028e/src/wheel/bdist_wheel.py#LL381C9-L381C85 + if not Path(self.dist_dir).exists(): + Path(self.dist_dir).mkdir(parents=True, exist_ok=True) + + impl_tag, abi_tag, plat_tag = self.get_tag() + archive_basename = f"{self.wheel_dist_name}-{impl_tag}-{abi_tag}-{plat_tag}" + + wheel_path = Path(self.dist_dir) / (archive_basename + ".whl") + print(f"Raw wheel path: {wheel_path}") + Path(wheel_filename).rename(wheel_path) + except (urllib.error.HTTPError, urllib.error.URLError): + print("Precompiled wheel not found. Building from source...") + # If the wheel could not be downloaded, build from source + super().run() + + +class NinjaBuildExtension(BuildExtension): + """ + Custom build extension that tells Ninja how many jobs to run. + + Credit: https://github.com/Dao-AILab/flash-attention/blob/main/setup.py + """ + + def __init__(self, *args: list[Any], **kwargs: dict[str, Any]) -> None: + # do not override env MAX_JOBS if already exists + if not os.environ.get("MAX_JOBS"): + try: + import psutil + + # calculate the maximum allowed NUM_JOBS based on cores + max_num_jobs_cores = max(1, os.cpu_count() // 2) + + # calculate the maximum allowed NUM_JOBS based on free memory + free_memory_gb = psutil.virtual_memory().available / ( + 1024**3 + ) # free memory in GB + max_num_jobs_memory = int( + free_memory_gb / 9, + ) # each JOB peak memory cost is ~8-9GB when threads = 4 + + # pick lower value of jobs based on cores vs memory metric to minimize + # oom and swap usage during compilation + max_jobs = max(1, min(max_num_jobs_cores, max_num_jobs_memory)) + os.environ["MAX_JOBS"] = str(max_jobs) + except ImportError: + warnings.warn( + "psutil not found, install psutil and ninja to get better build " + "performance", + stacklevel=1, + ) + + super().__init__(*args, **kwargs) + + +if SKIP_CUDA_BUILD: + warnings.warn( + "SKIP_CUDA_BUILD is set to 1, installing fouroversix without quantization and " + "matmul kernels", + stacklevel=1, + ) + + ext_modules = None +else: + if Path(".git").exists(): + subprocess.run( + [ + "git", + "submodule", + "update", + "--init", + "third_party/cutlass", + ], + check=True, + ) + elif not Path("third_party/cutlass").exists(): + msg = ( + "third_party/cutlass is missing, please use source distribution or git " + "clone" + ) + raise RuntimeError(msg) + + # The compiler flag -D_GLIBCXX_USE_CXX11_ABI is set to be the same as + # torch._C._GLIBCXX_USE_CXX11_ABI + # https://github.com/pytorch/pytorch/blob/8472c24e3b5b60150096486616d98b7bea01500b/torch/utils/cpp_extension.py#L920 + if FORCE_CXX11_ABI: + torch._C._GLIBCXX_USE_CXX11_ABI = True # noqa: SLF001 + + setup_dir = Path(__file__).parent + kernels_dir = setup_dir / "src" / "fouroversix" / "csrc" + sources = [ + path.relative_to(Path(__file__).parent).as_posix() + for ext in ["**/*.cu", "**/*.cpp"] + for path in kernels_dir.glob(ext) + ] + + cxx_compile_args = ["-std=c++17"] + nvcc_compile_args = [ + "-std=c++17", + "--expt-relaxed-constexpr", + "-Xcompiler", + "-funroll-loops", + "-Xcompiler", + "-finline-functions", + *get_cuda_gencodes(), + ] + + if CUTLASS_DEBUG: + nvcc_compile_args.extend( + [ + "-O0", + "-DCUTLASS_DEBUG_TRACE_LEVEL=3", + "-DCUTLASS_DEBUG_ENABLE=1", + "-g", + ], + ) + else: + cxx_compile_args.extend(["-O3"]) + nvcc_compile_args.extend(["-O3", "-DNDEBUG"]) + + ext_modules = [ + CUDAExtension( + "fouroversix._C", + sources, + extra_compile_args={"cxx": cxx_compile_args, "nvcc": nvcc_compile_args}, + include_dirs=[ + setup_dir / "third_party/cutlass/examples/common", + setup_dir / "third_party/cutlass/include", + setup_dir / "third_party/cutlass/tools/util/include", + kernels_dir / "include", + ], + ), + ] + +setup( + name=PACKAGE_NAME, + version=PACKAGE_VERSION, + ext_modules=ext_modules, + cmdclass={ + "bdist_wheel": CachedWheelsCommand, + "build_ext": NinjaBuildExtension, + }, + include_package_data=True, +) diff --git a/fouroversix/src/fouroversix/__init__.py b/fouroversix/src/fouroversix/__init__.py new file mode 100644 index 0000000..6b3adff --- /dev/null +++ b/fouroversix/src/fouroversix/__init__.py @@ -0,0 +1,33 @@ +from importlib.metadata import version + +from .matmul import fp4_matmul +from .model import ( + FourOverSixLinear, + ModelQuantizationConfig, + ModuleQuantizationConfig, + QuantizedModule, + quantize_model, +) +from .quantize import QuantizationConfig, QuantizedTensor, quantize_to_fp4 +from .utils import DataType, MatmulBackend, QuantizeBackend, RoundStyle, ScaleRule +from .weight_conversions import WeightConversions + +__version__ = version("fouroversix") + +__all__ = [ + "DataType", + "FourOverSixLinear", + "MatmulBackend", + "ModelQuantizationConfig", + "ModuleQuantizationConfig", + "QuantizationConfig", + "QuantizeBackend", + "QuantizedModule", + "QuantizedTensor", + "RoundStyle", + "ScaleRule", + "WeightConversions", + "fp4_matmul", + "quantize_model", + "quantize_to_fp4", +] diff --git a/fouroversix/src/fouroversix/csrc/bindings.cpp b/fouroversix/src/fouroversix/csrc/bindings.cpp new file mode 100644 index 0000000..4e3cdeb --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/bindings.cpp @@ -0,0 +1,36 @@ +#include +#include + +extern "C" +{ + PyObject *PyInit__C(void) + { + static struct PyModuleDef module_def = { + PyModuleDef_HEAD_INIT, + "_C", /* name of module */ + NULL, /* module documentation, may be NULL */ + -1, /* size of per-interpreter state of the module, + or -1 if the module keeps state in global variables. */ + NULL, /* methods */ + NULL, + NULL, + NULL, + NULL, + }; + return PyModule_Create(&module_def); + } +} + +namespace fouroversix +{ + TORCH_LIBRARY(fouroversix, m) + { + m.def("quantize_to_fp4(Tensor x, bool is_nvfp4, bool is_rtn, bool is_rht, bool is_2d, bool is_transpose, int selection_rule, int rbits) -> (Tensor, Tensor, Tensor)"); + m.def("gemm_mxfp4mxfp4_accum_fp32_out_bf16_tnt(Tensor A, Tensor B, Tensor A_sf, Tensor B_sf, Tensor alpha) -> Tensor"); + m.def("gemm_mxfp4mxfp4_accum_fp32_out_bf16_tnt_sm120(Tensor A, Tensor B, Tensor A_sf, Tensor B_sf, Tensor alpha) -> Tensor"); + m.def("gemm_nvfp4nvfp4_accum_fp32_out_bf16_tnt(Tensor A, Tensor B, Tensor A_sf, Tensor B_sf, Tensor alpha) -> Tensor"); + m.def("gemm_nvfp4nvfp4_accum_fp32_out_bf16_tnt_sm120(Tensor A, Tensor B, Tensor A_sf, Tensor B_sf, Tensor alpha) -> Tensor"); + m.def("gemm_nvfp4nvfp4_accum_fp32_out_fp16_tnt(Tensor A, Tensor B, Tensor A_sf, Tensor B_sf, Tensor alpha) -> Tensor"); + m.def("gemm_nvfp4nvfp4_accum_fp32_out_fp16_tnt_sm120(Tensor A, Tensor B, Tensor A_sf, Tensor B_sf, Tensor alpha) -> Tensor"); + } +} \ No newline at end of file diff --git a/fouroversix/src/fouroversix/csrc/fp4_gemm.cu b/fouroversix/src/fouroversix/csrc/fp4_gemm.cu new file mode 100644 index 0000000..695571b --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/fp4_gemm.cu @@ -0,0 +1,187 @@ +#include "cutlass/cutlass.h" + +#include "cute/tensor.hpp" +#include "cutlass/tensor_ref.h" +#include "cutlass/epilogue/thread/linear_combination.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/detail/sm100_blockscaled_layout.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/kernel/tile_scheduler_params.h" + +#include "cutlass/util/command_line.h" +#include "cutlass/util/distribution.h" +#include "cutlass/util/host_tensor.h" +#include "cutlass/util/packed_stride.hpp" +#include "cutlass/util/tensor_view_io.h" +#include "cutlass/util/reference/device/gemm.h" +#include "cutlass/util/reference/device/tensor_compare.h" +#include "cutlass/util/reference/host/tensor_fill.h" +#include "cutlass/util/reference/host/gett.hpp" +#include "cutlass/util/reference/host/tensor_norm.h" +#include "cutlass/util/reference/host/tensor_compare.h" + +#include "helper.h" + +#include +#include +#include +#include + +#include "element_traits.hpp" + +namespace fouroversix +{ + using namespace cute; + + // Adapted from example 72b + template , + typename ClusterShape = Shape<_2, _4, _1>, + typename KernelMainloopPolicy = cutlass::gemm::collective::KernelScheduleAuto, + int AlignmentA = 32, + typename ElementB = ElementA, + int AlignmentB = 32, + typename LayoutATag = cutlass::layout::RowMajor, + typename LayoutBTag = cutlass::layout::ColumnMajor, + typename ElementD = cutlass::bfloat16_t, + typename ArchTag = cutlass::arch::Sm100> + torch::Tensor gemm_fp4fp4_accum_fp32(torch::Tensor const &A, torch::Tensor const &B, torch::Tensor const &A_sf, torch::Tensor const &B_sf, torch::Tensor const &alpha) + { + at::cuda::CUDAGuard device_guard(A.device()); + + // C/D matrix configuration + using ElementC = void; // Element type for C matrix operand + using LayoutCTag = cutlass::layout::RowMajor; // Layout type for C matrix operand + using LayoutDTag = cutlass::layout::RowMajor; // Layout type for D matrix operand + + constexpr int AlignmentC = 1; // Memory access granularity/alignment of C matrix in units of elements (up to 16 bytes) + constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of D matrix in units of elements (up to 16 bytes) + + // Kernel functional config + using ElementAccumulator = float; // Element type for internal accumulation + using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; // Operator class tag + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementAccumulator, + ElementC, LayoutCTag, AlignmentC, + ElementD, LayoutDTag, AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto // Epilogue schedule policy + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + ElementA, LayoutATag, AlignmentA, + ElementB, LayoutBTag, AlignmentB, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + KernelMainloopPolicy>::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, // Indicates ProblemShape + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + // Reference device GEMM implementation type + using StrideA = typename Gemm::GemmKernel::StrideA; + using StrideB = typename Gemm::GemmKernel::StrideB; + using StrideC = typename Gemm::GemmKernel::StrideC; + using StrideD = typename Gemm::GemmKernel::StrideD; + + using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA; // Scale Factor tensors have an interleaved layout. Bring Layout instead of stride. + using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB; // Scale Factor tensors have an interleaved layout. Bring Layout instead of stride. + + using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + + torch::checkAllContiguous("gemm_fp4fp4_accum_fp32_out_bf16", {{A, "A", 0}, {B, "B", 1}, {A_sf, "A_sf", 2}, {B_sf, "B_sf", 3}, {alpha, "alpha", 4}}); + torch::checkDeviceType("gemm_fp4fp4_accum_fp32_out_bf16", {A, B, A_sf, B_sf, alpha}, at::DeviceType::CUDA); + torch::checkAllSameGPU("gemm_fp4fp4_accum_fp32_out_bf16", {{A, "A", 0}, {B, "B", 1}, {A_sf, "A_sf", 2}, {B_sf, "B_sf", 3}, {alpha, "alpha", 4}}); + + check_block_scale_factor_type(A_sf, "A_sf"); + check_block_scale_factor_type(B_sf, "B_sf"); + + auto [M, N, K] = check_and_get_fp4_matmul_dims(A, B, A_sf, B_sf); + auto D = torch::empty({M, N}, torch::dtype(element_traits::scalar_type).device(A.device())); + + Gemm gemm; + + // Create stride and layout information for the packed tensors + // For packed NVFP4 tensors, we need to use the appropriate stride and layout + StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, {M, K, 1}); + StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, {N, K, 1}); + StrideC stride_C = cutlass::make_cute_packed_stride(StrideC{}, {M, N, 1}); + StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, {M, N, 1}); + + // Create scale factor layouts + LayoutSFA layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(make_shape(M, N, K, 1)); + LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(make_shape(M, N, K, 1)); + + typename Gemm::Arguments args{ + cutlass::gemm::GemmUniversalMode::kGemm, + {M, N, K, 1}, + {// Mainloop arguments + static_cast(A.data_ptr()), stride_A, + static_cast(B.data_ptr()), stride_B, + static_cast(A_sf.data_ptr()), layout_SFA, + static_cast(B_sf.data_ptr()), layout_SFB}, + {// Epilogue arguments + {1.0f, 0.0f}, + nullptr, + stride_C, + static_cast(D.data_ptr()), + stride_D}}; + + args.epilogue.thread.alpha_ptr = static_cast(alpha.data_ptr()); + + // Check if the problem size is supported or not + CUTLASS_CHECK(gemm.can_implement(args)); + + // Initialize CUTLASS kernel with arguments and workspace pointer + CUTLASS_CHECK(gemm.initialize(args)); + + CUTLASS_CHECK(gemm.run(args)); + + return D; + } + + TORCH_LIBRARY_IMPL(fouroversix, CUDA, m) + { + /* MXFP4 */ + m.impl("gemm_mxfp4mxfp4_accum_fp32_out_bf16_tnt", + &gemm_fp4fp4_accum_fp32< + cutlass::mx_float4_t, + Shape<_256, _256, _256>, + Shape<_4, _1, _1>, + cutlass::gemm::KernelTmaWarpSpecialized2SmMxf4Sm100>); + + /* NVFP4 */ + m.impl("gemm_nvfp4nvfp4_accum_fp32_out_bf16_tnt", + &gemm_fp4fp4_accum_fp32< + cutlass::nv_float4_t, + Shape<_256, _256, _256>, + Shape<_4, _1, _1>, + cutlass::gemm::KernelTmaWarpSpecialized2SmNvf4Sm100>); + + m.impl("gemm_nvfp4nvfp4_accum_fp32_out_fp16_tnt", + &gemm_fp4fp4_accum_fp32< + cutlass::nv_float4_t, + Shape<_256, _256, _256>, + Shape<_4, _1, _1>, + cutlass::gemm::KernelTmaWarpSpecialized2SmNvf4Sm100, + 32, + cutlass::nv_float4_t, + 32, + cutlass::layout::RowMajor, + cutlass::layout::ColumnMajor, + cutlass::half_t>); + } +} \ No newline at end of file diff --git a/fouroversix/src/fouroversix/csrc/fp4_gemm_sm120.cu b/fouroversix/src/fouroversix/csrc/fp4_gemm_sm120.cu new file mode 100644 index 0000000..a828c5f --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/fp4_gemm_sm120.cu @@ -0,0 +1,414 @@ +#include "cutlass/cutlass.h" + +#include "cute/tensor.hpp" +#include "cutlass/tensor_ref.h" +#include "cutlass/epilogue/thread/linear_combination.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/detail/sm100_blockscaled_layout.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/kernel/tile_scheduler_params.h" + +#include "cutlass/util/command_line.h" +#include "cutlass/util/distribution.h" +#include "cutlass/util/host_tensor.h" +#include "cutlass/util/packed_stride.hpp" +#include "cutlass/util/tensor_view_io.h" +#include "cutlass/util/reference/device/gemm.h" +#include "cutlass/util/reference/device/tensor_compare.h" +#include "cutlass/util/reference/host/tensor_fill.h" +#include "cutlass/util/reference/host/gett.hpp" +#include "cutlass/util/reference/host/tensor_norm.h" +#include "cutlass/util/reference/host/tensor_compare.h" + +#include "helper.h" + +#include +#include +#include +#include + +#include "element_traits.hpp" + +// NOTE: There seems to be an issue with NVCC that causes runtime errors in SM120 GEMMs +// when they are compiled as templated functions. As a result, this file contains only +// non-templated functions. See: https://github.com/NVIDIA/cutlass/issues/2478 +namespace fouroversix +{ + using namespace cute; + + // Adapted from example 72b + torch::Tensor gemm_mxfp4mxfp4_accum_fp32_out_bf16_tnt_sm120(torch::Tensor const &A, torch::Tensor const &B, torch::Tensor const &A_sf, torch::Tensor const &B_sf, torch::Tensor const &alpha) + { + at::cuda::CUDAGuard device_guard(A.device()); + + using ElementA = cutlass::mx_float4_t; + using ElementB = cutlass::mx_float4_t; + using ElementD = cutlass::bfloat16_t; + + using MmaTileShape = Shape<_128, _128, _128>; + using ClusterShape = Shape<_1, _1, _1>; + + // C/D matrix configuration + using ElementC = void; // Element type for C matrix operand + using LayoutATag = cutlass::layout::RowMajor; // Layout type for A matrix operand + using LayoutBTag = cutlass::layout::ColumnMajor; // Layout type for B matrix operand + using LayoutCTag = cutlass::layout::RowMajor; // Layout type for C matrix operand + using LayoutDTag = cutlass::layout::RowMajor; // Layout type for D matrix operand + + constexpr int AlignmentA = 32; + constexpr int AlignmentB = 32; + constexpr int AlignmentC = 1; // Memory access granularity/alignment of C matrix in units of elements (up to 16 bytes) + constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of D matrix in units of elements (up to 16 bytes) + + // Kernel functional config + using ElementAccumulator = float; // Element type for internal accumulation + using ArchTag = cutlass::arch::Sm120; + using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; // Operator class tag + using KernelMainloopPolicy = cutlass::gemm::KernelTmaWarpSpecializedCooperative; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementAccumulator, + ElementC, LayoutCTag, AlignmentC, + ElementD, LayoutDTag, AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto // Epilogue schedule policy + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + ElementA, LayoutATag, AlignmentA, + ElementB, LayoutBTag, AlignmentB, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + KernelMainloopPolicy>::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, // Indicates ProblemShape + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + // Reference device GEMM implementation type + using StrideA = typename Gemm::GemmKernel::StrideA; + using StrideB = typename Gemm::GemmKernel::StrideB; + using StrideC = typename Gemm::GemmKernel::StrideC; + using StrideD = typename Gemm::GemmKernel::StrideD; + + using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA; // Scale Factor tensors have an interleaved layout. Bring Layout instead of stride. + using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB; // Scale Factor tensors have an interleaved layout. Bring Layout instead of stride. + + using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + + torch::checkAllContiguous("gemm_fp4fp4_accum_fp32_out_bf16", {{A, "A", 0}, {B, "B", 1}, {A_sf, "A_sf", 2}, {B_sf, "B_sf", 3}, {alpha, "alpha", 4}}); + torch::checkDeviceType("gemm_fp4fp4_accum_fp32_out_bf16", {A, B, A_sf, B_sf, alpha}, at::DeviceType::CUDA); + torch::checkAllSameGPU("gemm_fp4fp4_accum_fp32_out_bf16", {{A, "A", 0}, {B, "B", 1}, {A_sf, "A_sf", 2}, {B_sf, "B_sf", 3}, {alpha, "alpha", 4}}); + + check_block_scale_factor_type(A_sf, "A_sf"); + check_block_scale_factor_type(B_sf, "B_sf"); + + auto [M, N, K] = check_and_get_fp4_matmul_dims(A, B, A_sf, B_sf); + auto D = torch::empty({M, N}, torch::dtype(element_traits::scalar_type).device(A.device())); + + Gemm gemm; + + // Create stride and layout information for the packed tensors + // For packed NVFP4 tensors, we need to use the appropriate stride and layout + StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, {M, K, 1}); + StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, {N, K, 1}); + StrideC stride_C = cutlass::make_cute_packed_stride(StrideC{}, {M, N, 1}); + StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, {M, N, 1}); + + // Create scale factor layouts + LayoutSFA layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(make_shape(M, N, K, 1)); + LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(make_shape(M, N, K, 1)); + + typename Gemm::Arguments args{ + cutlass::gemm::GemmUniversalMode::kGemm, + {M, N, K, 1}, + {// Mainloop arguments + static_cast(A.data_ptr()), stride_A, + static_cast(B.data_ptr()), stride_B, + static_cast(A_sf.data_ptr()), layout_SFA, + static_cast(B_sf.data_ptr()), layout_SFB}, + {// Epilogue arguments + {1.0f, 0.0f}, + nullptr, + stride_C, + static_cast(D.data_ptr()), + stride_D}}; + + args.epilogue.thread.alpha_ptr = static_cast(alpha.data_ptr()); + + // Check if the problem size is supported or not + CUTLASS_CHECK(gemm.can_implement(args)); + + // Initialize CUTLASS kernel with arguments and workspace pointer + CUTLASS_CHECK(gemm.initialize(args)); + + CUTLASS_CHECK(gemm.run(args)); + + return D; + } + + // Adapted from example 72b + torch::Tensor gemm_nvfp4nvfp4_accum_fp32_out_bf16_tnt_sm120(torch::Tensor const &A, torch::Tensor const &B, torch::Tensor const &A_sf, torch::Tensor const &B_sf, torch::Tensor const &alpha) + { + at::cuda::CUDAGuard device_guard(A.device()); + + using ElementA = cutlass::nv_float4_t; + using ElementB = cutlass::nv_float4_t; + using ElementD = cutlass::bfloat16_t; + + using MmaTileShape = Shape<_128, _128, _128>; + using ClusterShape = Shape<_1, _1, _1>; + + // C/D matrix configuration + using ElementC = void; // Element type for C matrix operand + using LayoutATag = cutlass::layout::RowMajor; // Layout type for A matrix operand + using LayoutBTag = cutlass::layout::ColumnMajor; // Layout type for B matrix operand + using LayoutCTag = cutlass::layout::RowMajor; // Layout type for C matrix operand + using LayoutDTag = cutlass::layout::RowMajor; // Layout type for D matrix operand + + constexpr int AlignmentA = 32; + constexpr int AlignmentB = 32; + constexpr int AlignmentC = 1; // Memory access granularity/alignment of C matrix in units of elements (up to 16 bytes) + constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of D matrix in units of elements (up to 16 bytes) + + // Kernel functional config + using ElementAccumulator = float; // Element type for internal accumulation + using ArchTag = cutlass::arch::Sm120; + using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; // Operator class tag + using KernelMainloopPolicy = cutlass::gemm::KernelTmaWarpSpecializedCooperative; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementAccumulator, + ElementC, LayoutCTag, AlignmentC, + ElementD, LayoutDTag, AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto // Epilogue schedule policy + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + ElementA, LayoutATag, AlignmentA, + ElementB, LayoutBTag, AlignmentB, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + KernelMainloopPolicy>::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, // Indicates ProblemShape + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + // Reference device GEMM implementation type + using StrideA = typename Gemm::GemmKernel::StrideA; + using StrideB = typename Gemm::GemmKernel::StrideB; + using StrideC = typename Gemm::GemmKernel::StrideC; + using StrideD = typename Gemm::GemmKernel::StrideD; + + using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA; // Scale Factor tensors have an interleaved layout. Bring Layout instead of stride. + using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB; // Scale Factor tensors have an interleaved layout. Bring Layout instead of stride. + + using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + + torch::checkAllContiguous("gemm_fp4fp4_accum_fp32_out_bf16", {{A, "A", 0}, {B, "B", 1}, {A_sf, "A_sf", 2}, {B_sf, "B_sf", 3}, {alpha, "alpha", 4}}); + torch::checkDeviceType("gemm_fp4fp4_accum_fp32_out_bf16", {A, B, A_sf, B_sf, alpha}, at::DeviceType::CUDA); + torch::checkAllSameGPU("gemm_fp4fp4_accum_fp32_out_bf16", {{A, "A", 0}, {B, "B", 1}, {A_sf, "A_sf", 2}, {B_sf, "B_sf", 3}, {alpha, "alpha", 4}}); + + check_block_scale_factor_type(A_sf, "A_sf"); + check_block_scale_factor_type(B_sf, "B_sf"); + + auto [M, N, K] = check_and_get_fp4_matmul_dims(A, B, A_sf, B_sf); + auto D = torch::empty({M, N}, torch::dtype(element_traits::scalar_type).device(A.device())); + + Gemm gemm; + + // Create stride and layout information for the packed tensors + // For packed NVFP4 tensors, we need to use the appropriate stride and layout + StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, {M, K, 1}); + StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, {N, K, 1}); + StrideC stride_C = cutlass::make_cute_packed_stride(StrideC{}, {M, N, 1}); + StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, {M, N, 1}); + + // Create scale factor layouts + LayoutSFA layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(make_shape(M, N, K, 1)); + LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(make_shape(M, N, K, 1)); + + typename Gemm::Arguments args{ + cutlass::gemm::GemmUniversalMode::kGemm, + {M, N, K, 1}, + {// Mainloop arguments + static_cast(A.data_ptr()), stride_A, + static_cast(B.data_ptr()), stride_B, + static_cast(A_sf.data_ptr()), layout_SFA, + static_cast(B_sf.data_ptr()), layout_SFB}, + {// Epilogue arguments + {1.0f, 0.0f}, + nullptr, + stride_C, + static_cast(D.data_ptr()), + stride_D}}; + + args.epilogue.thread.alpha_ptr = static_cast(alpha.data_ptr()); + + // Check if the problem size is supported or not + CUTLASS_CHECK(gemm.can_implement(args)); + + // Initialize CUTLASS kernel with arguments and workspace pointer + CUTLASS_CHECK(gemm.initialize(args)); + + CUTLASS_CHECK(gemm.run(args)); + + return D; + } + + // Adapted from example 72b + torch::Tensor gemm_nvfp4nvfp4_accum_fp32_out_fp16_tnt_sm120(torch::Tensor const &A, torch::Tensor const &B, torch::Tensor const &A_sf, torch::Tensor const &B_sf, torch::Tensor const &alpha) + { + at::cuda::CUDAGuard device_guard(A.device()); + + using ElementA = cutlass::nv_float4_t; + using ElementB = cutlass::nv_float4_t; + using ElementD = cutlass::half_t; + + using MmaTileShape = Shape<_128, _128, _128>; + using ClusterShape = Shape<_1, _1, _1>; + + // C/D matrix configuration + using ElementC = void; // Element type for C matrix operand + using LayoutATag = cutlass::layout::RowMajor; // Layout type for A matrix operand + using LayoutBTag = cutlass::layout::ColumnMajor; // Layout type for B matrix operand + using LayoutCTag = cutlass::layout::RowMajor; // Layout type for C matrix operand + using LayoutDTag = cutlass::layout::RowMajor; // Layout type for D matrix operand + + constexpr int AlignmentA = 32; + constexpr int AlignmentB = 32; + constexpr int AlignmentC = 1; // Memory access granularity/alignment of C matrix in units of elements (up to 16 bytes) + constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of D matrix in units of elements (up to 16 bytes) + + // Kernel functional config + using ElementAccumulator = float; // Element type for internal accumulation + using ArchTag = cutlass::arch::Sm120; + using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; // Operator class tag + using KernelMainloopPolicy = cutlass::gemm::KernelTmaWarpSpecializedCooperative; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementAccumulator, + ElementC, LayoutCTag, AlignmentC, + ElementD, LayoutDTag, AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto // Epilogue schedule policy + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + ElementA, LayoutATag, AlignmentA, + ElementB, LayoutBTag, AlignmentB, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + KernelMainloopPolicy>::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, // Indicates ProblemShape + CollectiveMainloop, + CollectiveEpilogue, + void>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + // Reference device GEMM implementation type + using StrideA = typename Gemm::GemmKernel::StrideA; + using StrideB = typename Gemm::GemmKernel::StrideB; + using StrideC = typename Gemm::GemmKernel::StrideC; + using StrideD = typename Gemm::GemmKernel::StrideD; + + using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA; // Scale Factor tensors have an interleaved layout. Bring Layout instead of stride. + using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB; // Scale Factor tensors have an interleaved layout. Bring Layout instead of stride. + + using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + + torch::checkAllContiguous("gemm_fp4fp4_accum_fp32_out_bf16", {{A, "A", 0}, {B, "B", 1}, {A_sf, "A_sf", 2}, {B_sf, "B_sf", 3}, {alpha, "alpha", 4}}); + torch::checkDeviceType("gemm_fp4fp4_accum_fp32_out_bf16", {A, B, A_sf, B_sf, alpha}, at::DeviceType::CUDA); + torch::checkAllSameGPU("gemm_fp4fp4_accum_fp32_out_bf16", {{A, "A", 0}, {B, "B", 1}, {A_sf, "A_sf", 2}, {B_sf, "B_sf", 3}, {alpha, "alpha", 4}}); + + check_block_scale_factor_type(A_sf, "A_sf"); + check_block_scale_factor_type(B_sf, "B_sf"); + + auto [M, N, K] = check_and_get_fp4_matmul_dims(A, B, A_sf, B_sf); + auto D = torch::empty({M, N}, torch::dtype(element_traits::scalar_type).device(A.device())); + + Gemm gemm; + + // Create stride and layout information for the packed tensors + // For packed NVFP4 tensors, we need to use the appropriate stride and layout + StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, {M, K, 1}); + StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, {N, K, 1}); + StrideC stride_C = cutlass::make_cute_packed_stride(StrideC{}, {M, N, 1}); + StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, {M, N, 1}); + + // Create scale factor layouts + LayoutSFA layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(make_shape(M, N, K, 1)); + LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(make_shape(M, N, K, 1)); + + typename Gemm::Arguments args{ + cutlass::gemm::GemmUniversalMode::kGemm, + {M, N, K, 1}, + {// Mainloop arguments + static_cast(A.data_ptr()), stride_A, + static_cast(B.data_ptr()), stride_B, + static_cast(A_sf.data_ptr()), layout_SFA, + static_cast(B_sf.data_ptr()), layout_SFB}, + {// Epilogue arguments + {1.0f, 0.0f}, + nullptr, + stride_C, + static_cast(D.data_ptr()), + stride_D}}; + + args.epilogue.thread.alpha_ptr = static_cast(alpha.data_ptr()); + + // Check if the problem size is supported or not + CUTLASS_CHECK(gemm.can_implement(args)); + + // Initialize CUTLASS kernel with arguments and workspace pointer + CUTLASS_CHECK(gemm.initialize(args)); + + CUTLASS_CHECK(gemm.run(args)); + + return D; + } + + TORCH_LIBRARY_IMPL(fouroversix, CUDA, m) + { + // See the note at the top of the file for more information about why these + // functions are not templated and not included in fp4_gemm.cu. + + /* MXFP4 */ + m.impl("gemm_mxfp4mxfp4_accum_fp32_out_bf16_tnt_sm120", + &gemm_mxfp4mxfp4_accum_fp32_out_bf16_tnt_sm120); + + /* NVFP4 */ + m.impl("gemm_nvfp4nvfp4_accum_fp32_out_bf16_tnt_sm120", + &gemm_nvfp4nvfp4_accum_fp32_out_bf16_tnt_sm120); + m.impl("gemm_nvfp4nvfp4_accum_fp32_out_fp16_tnt_sm120", + &gemm_nvfp4nvfp4_accum_fp32_out_fp16_tnt_sm120); + } +} \ No newline at end of file diff --git a/fouroversix/src/fouroversix/csrc/include/element_traits.hpp b/fouroversix/src/fouroversix/csrc/include/element_traits.hpp new file mode 100644 index 0000000..91c0282 --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/include/element_traits.hpp @@ -0,0 +1,171 @@ +#include "cutlass/bfloat16.h" +#include "cutlass/float_subbyte.h" +#include "cutlass/half.h" +#include "cutlass/layout/matrix.h" +#include +#include + +namespace fouroversix +{ + template + struct element_traits; + + template <> + struct element_traits> + { + static constexpr size_t block_size = 32; + static constexpr size_t packing_factor = 2; + static constexpr size_t packed_bytes = 1; + static constexpr at::ScalarType block_scale_factor_type = at::kFloat8_e8m0fnu; + static constexpr at::ScalarType scalar_type = at::kByte; + }; + + template <> + struct element_traits> + { + static constexpr size_t block_size = 32; + static constexpr size_t packing_factor = 4; + static constexpr size_t packed_bytes = 3; + static constexpr at::ScalarType block_scale_factor_type = at::kFloat8_e8m0fnu; + static constexpr at::ScalarType scalar_type = at::kByte; + }; + + template <> + struct element_traits> + { + static constexpr size_t block_size = 32; + static constexpr size_t packing_factor = 1; + static constexpr size_t packed_bytes = 1; + static constexpr at::ScalarType block_scale_factor_type = at::kFloat8_e8m0fnu; + static constexpr at::ScalarType scalar_type = at::kByte; + }; + + template <> + struct element_traits> + { + static constexpr size_t block_size = 16; + static constexpr size_t packing_factor = 2; + static constexpr size_t packed_bytes = 1; + static constexpr at::ScalarType block_scale_factor_type = at::kFloat8_e4m3fn; + static constexpr at::ScalarType scalar_type = at::kByte; + }; + + template <> + struct element_traits + { + static constexpr at::ScalarType scalar_type = at::kBFloat16; + }; + + template <> + struct element_traits + { + static constexpr at::ScalarType scalar_type = at::kHalf; + }; + + template <> + struct element_traits + { + static constexpr at::ScalarType scalar_type = at::kFloat; + }; + + template + void check_block_scale_factor_type(const at::Tensor &t, const char *name) + { + TORCH_CHECK( + t.scalar_type() == element_traits::block_scale_factor_type, + name, " must be ", at::toString(element_traits::block_scale_factor_type)); + } + + template + std::tuple + check_and_get_bf16_matmul_dims(const at::Tensor &A, const at::Tensor &B) + { + TORCH_CHECK(A.scalar_type() == at::kBFloat16, "A must be bfloat16"); + TORCH_CHECK(B.scalar_type() == at::kBFloat16, "B must be bfloat16"); + TORCH_CHECK(A.dim() == 2, "A must be 2D"); + TORCH_CHECK(B.dim() == 2, "B must be 2D"); + + // Unpack sizes + int a_rows = A.size(0), a_cols = A.size(1); + int b_rows = B.size(0), b_cols = B.size(1); + + // Layout-based interpretation + int M, K_A; + if constexpr (std::is_same_v) + { + M = a_rows; + K_A = a_cols; + } + else if constexpr (std::is_same_v) + { + M = a_cols; + K_A = a_rows; + } + + int N, K_B; + if constexpr (std::is_same_v) + { + K_B = b_rows; + N = b_cols; + } + else if constexpr (std::is_same_v) + { + K_B = b_cols; + N = b_rows; + } + + TORCH_CHECK(K_A == K_B, "Inner dims mismatch: ", K_A, " vs ", K_B); + + return {M, N, K_A}; + } + + template + std::tuple + check_and_get_fp4_matmul_dims(const at::Tensor &A, const at::Tensor &B, + const at::Tensor &A_sf, const at::Tensor &B_sf) + { + TORCH_CHECK(A.scalar_type() == at::kByte, "A must be uint8"); + TORCH_CHECK(B.scalar_type() == at::kByte, "B must be uint8"); + TORCH_CHECK(A.dim() == 2, "A must be 2D"); + TORCH_CHECK(B.dim() == 2, "B must be 2D"); + TORCH_CHECK(A.size(1) >= 16, "A K-dim must be >= 16"); + TORCH_CHECK(B.size(1) >= 16, "B K-dim must be >= 16"); + + // Unpack 4-bit logical sizes + int a_rows = A.size(0), a_cols = A.size(1) * element_traits::packing_factor / element_traits::packed_bytes; + int b_rows = B.size(0), b_cols = B.size(1) * element_traits::packing_factor / element_traits::packed_bytes; + + // Layout-based interpretation + int M, K_A; + if constexpr (std::is_same_v) + { + M = a_rows; + K_A = a_cols; + } + else if constexpr (std::is_same_v) + { + M = a_cols; + K_A = a_rows; + } + + int N, K_B; + if constexpr (std::is_same_v) + { + K_B = b_rows; + N = b_cols; + } + else if constexpr (std::is_same_v) + { + K_B = b_cols; + N = b_rows; + } + + TORCH_CHECK(K_A == K_B, "Inner dims mismatch: ", K_A, " vs ", K_B); + + // Scale factor checks + TORCH_CHECK(A_sf.numel() * element_traits::block_size == size_t(a_rows) * size_t(a_cols), "A_sf size mismatch"); + TORCH_CHECK(B_sf.numel() * element_traits::block_size == size_t(b_rows) * size_t(b_cols), "B_sf size mismatch"); + + return {M, N, K_A}; + } +} \ No newline at end of file diff --git a/fouroversix/src/fouroversix/csrc/include/fp4_quant.h b/fouroversix/src/fouroversix/csrc/include/fp4_quant.h new file mode 100644 index 0000000..b4b5145 --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/include/fp4_quant.h @@ -0,0 +1,57 @@ +/****************************************************************************** + * Copyright (c) 2025, FourOverSix Team. + ******************************************************************************/ + +#pragma once +#include + +namespace fouroversix +{ + + enum AdaptiveBlockScalingRuleType + { + STATIC_6 = 0, + STATIC_4 = 1, + MAE_4o6 = 2, + MSE_4o6 = 3, + ABS_MAX_4o6 = 4, + }; + + struct FP4_quant_params + { + using index_t = int64_t; + void *__restrict__ x_ptr; + void *__restrict__ x_rht_ptr; + void *__restrict__ x_e2m1_ptr; + void *__restrict__ x_sf_ptr; + void *__restrict__ x_sft_ptr; + void *__restrict__ amax_ptr; + + int x_row_stride; + int x_col_stride; + int x_rht_row_stride; + int x_rht_col_stride; + int x_e2m1_row_stride; + int x_e2m1_col_stride; + int x_sf_row_stride; + int x_sf_col_stride; + int x_sft_row_stride; + int x_sft_col_stride; + + // The dimensions. + int M, N, M_rounded, N_rounded, M_sf, N_sf; + bool is_bf16; + bool is_nvfp4; + bool is_rtn; + bool is_rht; + bool is_4o6; + bool is_2d; + bool is_transpose; + int selection_rule; // 0: static_6, 1: static_4, 2: 4o6_mae, 3: 4o6_mse + int rbits; + }; + + template + void run_fp4_quant_(FP4_quant_params ¶ms, cudaStream_t stream); + +} // namespace fouroversix diff --git a/fouroversix/src/fouroversix/csrc/include/fp4_quant_kernel.h b/fouroversix/src/fouroversix/csrc/include/fp4_quant_kernel.h new file mode 100644 index 0000000..4cfe797 --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/include/fp4_quant_kernel.h @@ -0,0 +1,505 @@ +/****************************************************************************** + * Copyright (c) 2024, Tri Dao. + * Adapted by Junxian Guo from https://github.com/Dao-AILab/flash-attention/blob/main/csrc/flash_attn/src/flash_fwd_kernel.h + * Copyright (c) 2025, FourOverSix Team. + ******************************************************************************/ + +#pragma once + +// #include "philox_unpack.cuh" // For at::cuda::philox::unpack + +#include +#include + +#include +#include +#include + +#include "kernel_traits.h" +#include "utils.h" +#include "hadamard_transform.h" +// #include "softmax.h" +// #include "mask.h" +// #include "dropout.h" +// #include "rotary.h" + +namespace fouroversix +{ + + using namespace cute; + + template + inline __device__ void compute_fp4_quant_prologue_block(const Params ¶ms, const int m_block, const int n_block) + { + // Type aliases + using Element = typename Kernel_traits::Element; + using ScaleFactor = typename Kernel_traits::ScaleFactor; + using index_t = typename Kernel_traits::index_t; + + // Compile-time constants + constexpr int kGroupN = Kernel_traits::kGroupN; + constexpr int kBlockM = Kernel_traits::kBlockM; + constexpr int kBlockN = Kernel_traits::kBlockN; + constexpr int kNWarps = Kernel_traits::kNWarps; + constexpr int kNumGroupsInRow = Kernel_traits::kNumGroupsInRow; + constexpr int kNumGroupsInCol = Kernel_traits::kNumGroupsInCol; + constexpr float E4M3_MAX_VALUE = Kernel_traits::E4M3_MAX_VALUE; + constexpr float E2M1_MAX_VALUE = Kernel_traits::E2M1_MAX_VALUE; + + constexpr AdaptiveBlockScalingRuleType kRule = static_cast(kSelectionRule); + constexpr bool Is_4o6 = kRule == AdaptiveBlockScalingRuleType::MAE_4o6 || + kRule == AdaptiveBlockScalingRuleType::MSE_4o6 || + kRule == AdaptiveBlockScalingRuleType::ABS_MAX_4o6; + + using VecTypeX = cutlass::Array; + using VecTypeXFloat = cutlass::Array; + using VecTypeSFT = cutlass::Array; + constexpr int kVecSizeSFT = 4; + + // Shared memory + extern __shared__ char smem[]; + + // Runtime variables + const int tidx = threadIdx.x; + const int num_groups = kNumGroupsInRow * kBlockM; + float *amax_ptr = reinterpret_cast(params.amax_ptr); + + // ------------------------------------------------------------------------- + // Tensor Definitions + // ------------------------------------------------------------------------- + + // Input X (Global Memory) + Tensor mX = make_tensor(make_gmem_ptr(reinterpret_cast(params.x_ptr)), + make_shape(params.M, params.N), + make_stride(params.x_row_stride, _1{})); + Tensor gX = local_tile(mX(_, _), Shape, Int>{}, + make_coord(m_block, n_block)); + + Tensor mXRHT = make_tensor(make_gmem_ptr(reinterpret_cast(params.x_rht_ptr)), + make_shape(params.M_rounded, params.N_rounded), + make_stride(params.x_rht_row_stride, _1{})); + Tensor gXRHT = local_tile(mXRHT(_, _), Shape, Int>{}, + make_coord(m_block, n_block)); + + // Scale Factor Temp SFT (Global Memory) + Tensor mSFT = make_tensor(make_gmem_ptr(reinterpret_cast(params.x_sft_ptr)), + make_shape(params.M, params.N_rounded / kGroupN), + make_stride(params.x_sft_row_stride, _1{})); + Tensor gSFT = local_tile(mSFT(_, _), Shape, Int>{}, + make_coord(m_block, n_block)); + + // Shared Memory Tensors + Tensor sX = make_tensor(make_smem_ptr(reinterpret_cast(smem)), + typename Kernel_traits::SmemLayoutX{}); + + // SFT in Shared Memory (placed after X) + Tensor sSFT = make_tensor(make_smem_ptr(reinterpret_cast(reinterpret_cast(sX.data().get()) + sizeof(Element) * size(sX))), + typename Kernel_traits::SmemLayoutSFT{}); + + // ------------------------------------------------------------------------- + // Data Loading (X -> Shared) + // ------------------------------------------------------------------------- + + typename Kernel_traits::GmemTiledCopyX gmem_tiled_copy_X; + auto gmem_thr_copy_X = gmem_tiled_copy_X.get_thread_slice(tidx); + + Tensor tXgX = gmem_thr_copy_X.partition_S(gX); + Tensor tXsX = gmem_thr_copy_X.partition_D(sX); + + // Construct predicates for bounds checking + Tensor cX = make_identity_tensor(make_shape(size<0>(sX), size<1>(sX))); + Tensor tXcX = gmem_thr_copy_X.partition_S(cX); + Tensor tXpX = make_tensor(make_shape(size<2>(tXcX))); + + for (int i = 0; i < size(tXpX); ++i) + { + tXpX(i) = get<1>(tXcX(0, 0, i)) < params.N - n_block * kBlockN; + } + + __syncthreads(); + + // Async copy from Global to Shared + fouroversix::copy( + gmem_tiled_copy_X, tXgX, tXsX, tXcX, tXpX, params.M - m_block * kBlockM); + + cute::cp_async_fence(); + fouroversix::cp_async_wait<0>(); + __syncthreads(); + + // ------------------------------------------------------------------------- + // Scale Factor Computation + // ------------------------------------------------------------------------- + + float thr_max = 0.0f; + for (int g_idx = tidx; g_idx < num_groups; g_idx += blockDim.x) + { + const int g_row = g_idx / kNumGroupsInRow; + const int g_col = g_idx % kNumGroupsInRow; + + VecTypeXFloat x_vec_float; + for (int i = 0; i < kGroupN; ++i) + { + x_vec_float[i] = static_cast(sX(g_row, g_col * kGroupN + i)); + } + if constexpr (Is_rht) + { + hadamard_quant_group(&x_vec_float[0]); + VecTypeX x_vec; +#pragma unroll + for (int i = 0; i < kGroupN; ++i) + { + // sX(g_row, g_col * kGroupN + i) = static_cast(x_vec_float[i]); + x_vec[i] = static_cast(x_vec_float[i]); + } + + *reinterpret_cast(&gXRHT(g_row, g_col * kGroupN)) = *reinterpret_cast(&x_vec); + } + // VecTypeX x_vec = *reinterpret_cast(&sX(g_row, g_col * kGroupN)); + + // Compute max absolute value in group + float sf = 0.0f; +#pragma unroll + for (int i = 0; i < kGroupN; ++i) + { + sf = max(sf, abs(x_vec_float[i])); + } + + thr_max = max(thr_max, sf); + sSFT(g_row, g_col) = sf; + } + + if constexpr (Is_2d) + { + __syncthreads(); + } + + if constexpr (Is_2d) + { + MaxOp max_op; + for (int g_idx = tidx; g_idx < num_groups; g_idx += blockDim.x) + { + const int g_row = g_idx % kNumGroupsInCol; + const int g_col = g_idx / kNumGroupsInCol; + float sf = sSFT(g_row, g_col); + float blk_sf = Allreduce::run(sf, max_op); // kGroupN is 16 or 32 + sSFT(g_row, g_col) = blk_sf; + __syncthreads(); + } + } + + // ------------------------------------------------------------------------- + // Normalization Constant Reduction (Block-wide Max) + // ------------------------------------------------------------------------- + + // Warp-level reduction + MaxOp max_op; + float warp_max = Allreduce<32>::run(thr_max, max_op); + + // Block-level reduction via shared memory + float *sRed = reinterpret_cast(smem); + if (tidx % 32 == 0) + { + sRed[tidx / 32] = warp_max; + } + __syncthreads(); + + if (tidx == 0) + { + float blk_max = 0.0f; +#pragma unroll + for (int i = 0; i < kNWarps; ++i) + { + blk_max = max(blk_max, sRed[i]); + } + atomicMaxFloat(amax_ptr, blk_max); + } + + // ------------------------------------------------------------------------- + // Write Back SFT (Shared -> Global) + // ------------------------------------------------------------------------- + + for (int r_idx = tidx; r_idx < kBlockM; r_idx += blockDim.x) + { +#pragma unroll + for (int i = 0; i < int(kBlockN / kGroupN); i += kVecSizeSFT) + { + *reinterpret_cast(&gSFT(r_idx, i)) = *reinterpret_cast(&sSFT(r_idx, i)); + } + } + } + + template + inline __device__ void compute_fp4_quant_prologue(const Params ¶ms) + { + // TODO: Implement the fp4 quant kernel + const int m_block = blockIdx.x; + // The block index for the batch. + const int n_block = blockIdx.y; + + fouroversix::compute_fp4_quant_prologue_block(params, m_block, n_block); + } + + template + inline __device__ void compute_fp4_quant_block(const Params ¶ms, const int m_block, const int n_block) + { + // Type aliases + using Element = typename Kernel_traits::Element; + using ScaleFactor = typename Kernel_traits::ScaleFactor; + using index_t = typename Kernel_traits::index_t; + + // Compile-time constants + constexpr int kGroupN = Kernel_traits::kGroupN; + constexpr int kBlockM = Kernel_traits::kBlockM; + constexpr int kBlockN = Kernel_traits::kBlockN; + constexpr int kBlockMSF = Kernel_traits::kBlockMSF; + constexpr int kBlockNSF = Kernel_traits::kBlockNSF; + constexpr int kNWarps = Kernel_traits::kNWarps; + constexpr int kNumGroupsInRow = Kernel_traits::kNumGroupsInRow; + constexpr int kNumGroupsInCol = Kernel_traits::kNumGroupsInCol; + constexpr float E4M3_MAX_VALUE = Kernel_traits::E4M3_MAX_VALUE; + + constexpr AdaptiveBlockScalingRuleType kRule = static_cast(kSelectionRule); + constexpr bool Is_4o6 = kRule == AdaptiveBlockScalingRuleType::MAE_4o6 || + kRule == AdaptiveBlockScalingRuleType::MSE_4o6 || + kRule == AdaptiveBlockScalingRuleType::ABS_MAX_4o6; + constexpr float E4M3_SCALE_4 = Is_4o6 ? Kernel_traits::E4M3_MAX_FOUROVERSIX : E4M3_MAX_VALUE; + constexpr float E4M3_SCALE_6 = Is_4o6 ? Kernel_traits::E4M3_MAX_FOUROVERSIX : E4M3_MAX_VALUE; + constexpr float E2M1_SCALE_4 = Is_4o6 ? 6.0f : 4.0f; + constexpr float E2M1_SCALE_6 = 6.0f; + + constexpr int kSmemBlockInRow = int(kNumGroupsInRow / 4); + constexpr int kSmemBlockInCol = int(kBlockM / 128); + + using VecTypeXe2m1 = std::conditional_t, cutlass::Array>; + using VecTypeSFT = cutlass::Array; + using VecTypeSF = cutlass::Array; + using OutputType = cutlass::Array; + constexpr int kVecSizeXe2m1 = Is_nvfp4 ? 8 : 16; + constexpr int kVecSizeSFT = 4; + constexpr int kVecSizeSF = 16; + + // Shared memory + extern __shared__ char smem[]; + + // Runtime variables + const int tidx = threadIdx.x; + const int num_groups = kNumGroupsInRow * kBlockM; + // JXGuo: assure amax is not zero before calling this kernel + const float amax = *reinterpret_cast(params.amax_ptr); + + if (amax == 0.0f) + { + return; + } + + // ------------------------------------------------------------------------- + // Tensor Definitions + // ------------------------------------------------------------------------- + + // Input X (Global Memory) + // void *__restrict__ x_ptr = Is_rht ? params.x_rht_ptr : params.x_ptr; + Tensor mX = Is_rht ? make_tensor(make_gmem_ptr(reinterpret_cast(params.x_rht_ptr)), + make_shape(params.M_rounded, params.N_rounded), + make_stride(params.x_rht_row_stride, _1{})) + : make_tensor(make_gmem_ptr(reinterpret_cast(params.x_ptr)), + make_shape(params.M, params.N), + make_stride(params.x_row_stride, _1{})); + Tensor gX = local_tile(mX(_, _), Shape, Int>{}, + make_coord(m_block, n_block)); + + Tensor mXe2m1 = make_tensor(make_gmem_ptr(reinterpret_cast(params.x_e2m1_ptr)), + make_shape(params.M, params.N_rounded / 2), + make_stride(params.x_e2m1_row_stride, _1{})); + Tensor gXe2m1 = local_tile(mXe2m1(_, _), Shape, Int>{}, + make_coord(m_block, n_block)); + + // Scale Factor Temp SFT (Global Memory) + Tensor mSFT = make_tensor(make_gmem_ptr(reinterpret_cast(params.x_sft_ptr)), + make_shape(params.M, params.N_rounded / kGroupN), + make_stride(params.x_sft_row_stride, _1{})); + Tensor gSFT = local_tile(mSFT(_, _), Shape, Int>{}, + make_coord(m_block, n_block)); + + Tensor gSF = make_tensor(make_gmem_ptr(reinterpret_cast(params.x_sf_ptr)), + make_shape(params.M_sf, params.N_sf), + make_stride(params.x_sf_row_stride, _1{})); + // Tensor gSF = local_tile(mSF(_, _), Shape, Int<16>>{}, + // make_coord(m_block, n_block)); + + // Shared Memory Tensors + Tensor sX = make_tensor(make_smem_ptr(reinterpret_cast(smem)), + typename Kernel_traits::SmemLayoutX{}); + + // SFT in Shared Memory (placed after X) + Tensor sSFT = make_tensor(make_smem_ptr(reinterpret_cast(reinterpret_cast(sX.data().get()) + sizeof(Element) * size(sX))), + typename Kernel_traits::SmemLayoutSFT{}); + + Tensor sXe2m1 = make_tensor(make_smem_ptr(reinterpret_cast(reinterpret_cast(sSFT.data().get()) + sizeof(float) * size(sSFT))), + Shape, Int>{}, + Stride, _1>{}); + + Tensor sSF = make_tensor(make_smem_ptr(reinterpret_cast(reinterpret_cast(sXe2m1.data().get()) + sizeof(uint8_t) * size(sXe2m1))), + typename Kernel_traits::SmemLayoutSF{}); + + // ------------------------------------------------------------------------- + // Data Loading (X -> Shared) + // ------------------------------------------------------------------------- + + typename Kernel_traits::GmemTiledCopyX gmem_tiled_copy_X; + auto gmem_thr_copy_X = gmem_tiled_copy_X.get_thread_slice(tidx); + + Tensor tXgX = gmem_thr_copy_X.partition_S(gX); + Tensor tXsX = gmem_thr_copy_X.partition_D(sX); + + // Construct predicates for bounds checking + Tensor cX = make_identity_tensor(make_shape(size<0>(sX), size<1>(sX))); + Tensor tXcX = gmem_thr_copy_X.partition_S(cX); + Tensor tXpX = make_tensor(make_shape(size<2>(tXcX))); + + for (int i = 0; i < size(tXpX); ++i) + { + tXpX(i) = get<1>(tXcX(0, 0, i)) < params.N - n_block * kBlockN; + } + + __syncthreads(); + + // Async copy from Global to Shared + fouroversix::copy( + gmem_tiled_copy_X, tXgX, tXsX, tXcX, tXpX, params.M - m_block * kBlockM); + + cute::cp_async_fence(); + + // ------------------------------------------------------------------------- + // Data Loading (SFT -> Shared) + // ------------------------------------------------------------------------- + + for (int r_idx = tidx; r_idx < kBlockM; r_idx += blockDim.x) + { +#pragma unroll + for (int i = 0; i < int(kBlockN / kGroupN); i += kVecSizeSFT) + { + *reinterpret_cast(&sSFT(r_idx, i)) = *reinterpret_cast(&gSFT(r_idx, i)); + } + } + + fouroversix::cp_async_wait<0>(); + __syncthreads(); + + // ------------------------------------------------------------------------- + // Quantization + // ------------------------------------------------------------------------- + + for (int g_idx = tidx; g_idx < num_groups; g_idx += blockDim.x) + { + const int g_row = g_idx % kNumGroupsInCol; + const int g_col = g_idx / kNumGroupsInCol; + const float g_max = sSFT(g_row, g_col); + + const Tensor sGX = make_tensor(make_smem_ptr(sX.data() + g_row * kBlockN + g_col * kGroupN), + Shape, Int>{}, + Stride, _1>{}); + + OutputType res[int(kGroupN / 8)]; + float encode_scale; + float sf; + + if constexpr (Is_4o6) + { + encode_scale = E2M1_SCALE_6 * E4M3_SCALE_6 / amax; + + float sf_high_precision = g_max / E2M1_SCALE_6 * encode_scale; + float sf_[2] = {sf_high_precision * 1.5, sf_high_precision}; + + sf_[0] = static_cast(static_cast(sf_[0])); + sf_[1] = static_cast(static_cast(sf_[1])); + + sf = fp4_conversion(sGX, amax, sf_, res, params.rbits); + } + else + { + float sf_val = 0.0f; + if constexpr (kRule == AdaptiveBlockScalingRuleType::STATIC_6) + { + encode_scale = E4M3_SCALE_6 * E2M1_SCALE_6 / amax; + sf_val = clamp(g_max / E2M1_SCALE_6 * encode_scale, 0, E4M3_MAX_VALUE); + } + else if constexpr (kRule == AdaptiveBlockScalingRuleType::STATIC_4) + { + encode_scale = E2M1_SCALE_4 * E4M3_SCALE_4 / amax; + sf_val = clamp(g_max / E2M1_SCALE_4 * encode_scale, 0, E4M3_MAX_VALUE); + } + else + { + printf("in fp4_quant_block, kRule = %d, not supported\n", kRule); + assert(false); + } + + sf_val = static_cast(static_cast(sf_val)); + sf = fp4_conversion(sGX, amax, &sf_val, res, params.rbits); + } + + // Write quantized data + for (int i = 0; i < int(kGroupN / 8); ++i) + { + *reinterpret_cast(&sXe2m1(g_row, g_col * (kGroupN / 2) + i * 4)) = res[i]; + } + + // Write scale factor (layout: 128x4 blocks, 32 rows per block) + const int r_in_blk = g_row % 128; + const int c_in_blk = g_col % 4; + const int blk_row = int(g_row / 128); + const int blk_col = int(g_col / 4); + const int sf_row = 32 * (blk_row * kSmemBlockInRow + blk_col) + r_in_blk % 32; + const int sf_col = int(r_in_blk / 32) * 4 + c_in_blk; + sSF(sf_row, sf_col) = static_cast(sf); + __syncthreads(); + } + + // ------------------------------------------------------------------------- + // Write Back Xe2m1 (Shared -> Global) + // ------------------------------------------------------------------------- + + __syncthreads(); + + for (int r_idx = tidx; r_idx < kBlockM; r_idx += blockDim.x) + { +#pragma unroll + for (int i = 0; i < int(kBlockN / 2); i += kVecSizeXe2m1) + { + *reinterpret_cast(&gXe2m1(r_idx, i)) = *reinterpret_cast(&sXe2m1(r_idx, i)); + } + } + + // ------------------------------------------------------------------------- + // Write Back SF (Shared -> Global) + // ------------------------------------------------------------------------- + + const int gbl_blk_row_stride = int(params.N_rounded / (kGroupN * 4)); + const int gbl_blk_col_stride = 1; + const int gbl_blk_idx_base = (m_block * kSmemBlockInCol) * gbl_blk_row_stride + (n_block * kSmemBlockInRow) * gbl_blk_col_stride; + + static_assert(kVecSizeSF == kBlockNSF, "kVecSizeSF must be equal to kBlockNSF"); + for (int r_idx = tidx; r_idx < kBlockMSF; r_idx += blockDim.x) + { + const int loc_blk_idx = int(r_idx / 32); + const int loc_row = r_idx % 32; + const int loc_blk_row = int(loc_blk_idx / kSmemBlockInRow); + const int loc_blk_col = int(loc_blk_idx % kSmemBlockInRow); + const int gbl_blk_idx = gbl_blk_idx_base + loc_blk_row * gbl_blk_row_stride + loc_blk_col * gbl_blk_col_stride; + const index_t gbl_row = index_t(32) * gbl_blk_idx + loc_row; + *reinterpret_cast(&gSF(gbl_row, 0)) = *reinterpret_cast(&sSF(r_idx, 0)); + } + } + + template + inline __device__ void compute_fp4_quant(const Params ¶ms) + { + // TODO: Implement the fp4 quant kernel + const int m_block = blockIdx.x; + // The block index for the batch. + const int n_block = blockIdx.y; + + fouroversix::compute_fp4_quant_block(params, m_block, n_block); + } + +} // namespace fouroversix diff --git a/fouroversix/src/fouroversix/csrc/include/fp4_quant_launch_template.h b/fouroversix/src/fouroversix/csrc/include/fp4_quant_launch_template.h new file mode 100644 index 0000000..4168b3e --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/include/fp4_quant_launch_template.h @@ -0,0 +1,147 @@ +/****************************************************************************** + * Copyright (c) 2023, Tri Dao. + * Adapted by Junxian Guo from https://github.com/Dao-AILab/flash-attention/blob/main/csrc/flash_attn/src/flash_fwd_launch_template.h + * Copyright (c) 2025, FourOverSix Team. + ******************************************************************************/ + +#pragma once +#include // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK + +#include "static_switch.h" +#include "hardware_info.h" +#include "fp4_quant.h" +#include "fp4_quant_kernel.h" + +namespace fouroversix +{ + +// Determine if the architecture supports FLASH and define a macro to handle parameter modifiers +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 +#define ARCH_SUPPORTS_FLASH +#define KERNEL_PARAM_MODIFIER __grid_constant__ +#else +#define KERNEL_PARAM_MODIFIER +#endif + +// Define a macro for unsupported architecture handling to centralize the error message +#define FLASH_UNSUPPORTED_ARCH printf("FATAL: FourOverSix requires building with sm version sm100, but was built for < 10.0!"); + +// Use a macro to clean up kernel definitions +#define DEFINE_FP4_QUANT_KERNEL(kernelName, ...) \ + template \ + __global__ void kernelName(KERNEL_PARAM_MODIFIER const FP4_quant_params params) + + DEFINE_FP4_QUANT_KERNEL(fp4_quant_prologue_kernel, bool Is_nvfp4, bool Is_rht, bool Is_2d, bool Is_transpose, bool Is_rtn, int kSelectionRule) + { +#if defined(ARCH_SUPPORTS_FLASH) + fouroversix::compute_fp4_quant_prologue(params); +#else + FLASH_UNSUPPORTED_ARCH +#endif + } + + DEFINE_FP4_QUANT_KERNEL(fp4_quant_kernel, bool Is_nvfp4, bool Is_rht, bool Is_2d, bool Is_transpose, bool Is_rtn, int kSelectionRule) + { +#if defined(ARCH_SUPPORTS_FLASH) + fouroversix::compute_fp4_quant(params); +#else + FLASH_UNSUPPORTED_ARCH +#endif + } + + //////////////////////////////////////////////////////////////////////////////////////////////////// + + template + void launch_fp4_quant_prologue(FP4_quant_params ¶ms, cudaStream_t stream) + { + constexpr size_t smem_size = Kernel_traits::kSmemSize; + + const int num_m_block = (params.M + Kernel_traits::kBlockM - 1) / Kernel_traits::kBlockM; + const int num_n_block = (params.N + Kernel_traits::kBlockN - 1) / Kernel_traits::kBlockN; + dim3 grid(num_m_block, num_n_block); + BOOL_SWITCH(params.is_rtn, Is_rtn, [&] + { + BOOL_SWITCH(params.is_2d, Is_2d, [&] + { + SELECTION_RULE_SWITCH(params.selection_rule, kSelectionRule, [&] + { + auto kernel_prologue = &fp4_quant_prologue_kernel; + if (smem_size >= 48 * 1024) { + C10_CUDA_CHECK(cudaFuncSetAttribute( + kernel_prologue, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + } + kernel_prologue<<>>(params); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + }); + }); + } + + //////////////////////////////////////////////////////////////////////////////////////////////////// + + template + void launch_fp4_quant(FP4_quant_params ¶ms, cudaStream_t stream) + { + constexpr size_t smem_size = Kernel_traits::kSmemSize; + + const int num_m_block = (params.M + Kernel_traits::kBlockM - 1) / Kernel_traits::kBlockM; + const int num_n_block = (params.N + Kernel_traits::kBlockN - 1) / Kernel_traits::kBlockN; + dim3 grid(num_m_block, num_n_block); + BOOL_SWITCH(params.is_rtn, Is_rtn, [&] + { + BOOL_SWITCH(params.is_2d, Is_2d, [&] + { + SELECTION_RULE_SWITCH(params.selection_rule, kSelectionRule, [&] + { + auto kernel = &fp4_quant_kernel; + if (smem_size >= 48 * 1024) { + C10_CUDA_CHECK(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + } + kernel<<>>(params); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + }); + }); + } + + //////////////////////////////////////////////////////////////////////////////////////////////////// + // template> + + template + void run_mxfp4_quant(FP4_quant_params ¶ms, cudaStream_t stream) + { + constexpr bool Is_nvfp4 = false; + constexpr bool Is_rht = false; + launch_fp4_quant_prologue, Is_nvfp4, Is_rht, Is_transpose>(params, stream); + launch_fp4_quant, Is_nvfp4, Is_rht, Is_transpose>(params, stream); + } + + template + void run_mxfp4_quant_rht(FP4_quant_params ¶ms, cudaStream_t stream) + { + constexpr bool Is_nvfp4 = false; + constexpr bool Is_rht = true; + launch_fp4_quant_prologue, Is_nvfp4, Is_rht, Is_transpose>(params, stream); + launch_fp4_quant, Is_nvfp4, Is_rht, Is_transpose>(params, stream); + } + + template + void run_nvfp4_quant(FP4_quant_params ¶ms, cudaStream_t stream) + { + constexpr bool Is_nvfp4 = true; + constexpr bool Is_rht = false; + launch_fp4_quant_prologue, Is_nvfp4, Is_rht, Is_transpose>(params, stream); + launch_fp4_quant, Is_nvfp4, Is_rht, Is_transpose>(params, stream); + } + + template + void run_nvfp4_quant_rht(FP4_quant_params ¶ms, cudaStream_t stream) + { + constexpr bool Is_nvfp4 = true; + constexpr bool Is_rht = true; + launch_fp4_quant_prologue, Is_nvfp4, Is_rht, Is_transpose>(params, stream); + launch_fp4_quant, Is_nvfp4, Is_rht, Is_transpose>(params, stream); + } + +} // namespace fouroversix diff --git a/fouroversix/src/fouroversix/csrc/include/hadamard_transform.h b/fouroversix/src/fouroversix/csrc/include/hadamard_transform.h new file mode 100644 index 0000000..52ba23e --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/include/hadamard_transform.h @@ -0,0 +1,56 @@ + +/****************************************************************************** + * Copyright (c) 2023, Tri Dao. + * Adapted by Junxian Guo from https://github.com/Dao-AILab/fast-hadamard-transform/blob/master/csrc/code_gen.py + * Copyright (c) 2025, FourOverSix Team. + ******************************************************************************/ + +// This file is auto-generated. See "hadamard_code_gen.py" + + +#pragma once + + +namespace fouroversix { + +template +__device__ __forceinline__ void hadamard_quant_group_16(float x[16]) { + float out[16]; + out[0] = + x[0] + x[1] + x[2] - x[3] + x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] + x[13] - x[14] - x[15]; + out[1] = + x[0] - x[1] + x[2] + x[3] + x[4] + x[5] - x[6] + x[7] - x[8] + x[9] - x[10] - x[11] - x[12] - x[13] - x[14] + x[15]; + out[2] = + x[0] + x[1] - x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] + x[10] - x[11] - x[12] + x[13] + x[14] + x[15]; + out[3] = + x[0] - x[1] - x[2] - x[3] + x[4] + x[5] + x[6] - x[7] - x[8] + x[9] + x[10] + x[11] - x[12] - x[13] + x[14] - x[15]; + out[4] = + x[0] + x[1] + x[2] - x[3] - x[4] + x[5] + x[6] + x[7] - x[8] - x[9] - x[10] + x[11] + x[12] - x[13] + x[14] + x[15]; + out[5] = + x[0] - x[1] + x[2] + x[3] - x[4] - x[5] + x[6] - x[7] - x[8] + x[9] - x[10] - x[11] + x[12] + x[13] + x[14] - x[15]; + out[6] = + x[0] + x[1] - x[2] + x[3] - x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] - x[11] + x[12] - x[13] - x[14] - x[15]; + out[7] = + x[0] - x[1] - x[2] - x[3] - x[4] - x[5] - x[6] + x[7] - x[8] + x[9] + x[10] + x[11] + x[12] + x[13] - x[14] + x[15]; + out[8] = + x[0] + x[1] + x[2] - x[3] + x[4] - x[5] - x[6] - x[7] + x[8] + x[9] + x[10] - x[11] + x[12] - x[13] + x[14] + x[15]; + out[9] = + x[0] - x[1] + x[2] + x[3] + x[4] + x[5] - x[6] + x[7] + x[8] - x[9] + x[10] + x[11] + x[12] + x[13] + x[14] - x[15]; + out[10] = + x[0] + x[1] - x[2] + x[3] + x[4] - x[5] + x[6] + x[7] + x[8] + x[9] - x[10] + x[11] + x[12] - x[13] - x[14] - x[15]; + out[11] = + x[0] - x[1] - x[2] - x[3] + x[4] + x[5] + x[6] - x[7] + x[8] - x[9] - x[10] - x[11] + x[12] + x[13] - x[14] + x[15]; + out[12] = + x[0] + x[1] + x[2] - x[3] - x[4] + x[5] + x[6] + x[7] + x[8] + x[9] + x[10] - x[11] - x[12] + x[13] - x[14] - x[15]; + out[13] = + x[0] - x[1] + x[2] + x[3] - x[4] - x[5] + x[6] - x[7] + x[8] - x[9] + x[10] + x[11] - x[12] - x[13] - x[14] + x[15]; + out[14] = + x[0] + x[1] - x[2] + x[3] - x[4] + x[5] - x[6] - x[7] + x[8] + x[9] - x[10] + x[11] - x[12] + x[13] + x[14] + x[15]; + out[15] = + x[0] - x[1] - x[2] - x[3] - x[4] - x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11] - x[12] - x[13] + x[14] - x[15]; + #pragma unroll + for (int i = 0; i < 16; i++) { x[i] = static_cast(static_cast(out[i] / 4)); } +} + +template +__device__ __forceinline__ void hadamard_quant_group_32(float x[32]) { + hadamard_quant_group_16(x); + hadamard_quant_group_16(x + 16); +} + +template +__device__ __forceinline__ void hadamard_quant_group(float* x) { + if constexpr (Is_nvfp4) { + hadamard_quant_group_16(x); + } else { + hadamard_quant_group_32(x); + } +} + + +} // namespace fouroversix + diff --git a/fouroversix/src/fouroversix/csrc/include/hardware_info.h b/fouroversix/src/fouroversix/csrc/include/hardware_info.h new file mode 100644 index 0000000..b218a29 --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/include/hardware_info.h @@ -0,0 +1,41 @@ +/****************************************************************************** + * Copyright (c) 2024, Tri Dao. + ******************************************************************************/ + +#pragma once + +#include + +#if !defined(__CUDACC_RTC__) +#include "cuda_runtime.h" +#endif + +#define CHECK_CUDA(call) \ + do { \ + cudaError_t status_ = call; \ + if (status_ != cudaSuccess) { \ + fprintf(stderr, "CUDA error (%s:%d): %s\n", __FILE__, __LINE__, \ + cudaGetErrorString(status_)); \ + exit(1); \ + } \ + } while (0) + + +inline int get_current_device() { + int device; + CHECK_CUDA(cudaGetDevice(&device)); + return device; +} + +inline std::tuple get_compute_capability(int device) { + int capability_major, capability_minor; + CHECK_CUDA(cudaDeviceGetAttribute(&capability_major, cudaDevAttrComputeCapabilityMajor, device)); + CHECK_CUDA(cudaDeviceGetAttribute(&capability_minor, cudaDevAttrComputeCapabilityMinor, device)); + return {capability_major, capability_minor}; +} + +inline int get_num_sm(int device) { + int multiprocessor_count; + CHECK_CUDA(cudaDeviceGetAttribute(&multiprocessor_count, cudaDevAttrMultiProcessorCount, device)); + return multiprocessor_count; +} diff --git a/fouroversix/src/fouroversix/csrc/include/kernel_traits.h b/fouroversix/src/fouroversix/csrc/include/kernel_traits.h new file mode 100644 index 0000000..4f3d6e4 --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/include/kernel_traits.h @@ -0,0 +1,141 @@ +/****************************************************************************** + * Copyright (c) 2024, Tri Dao. + * Adapted by Junxian Guo from https://github.com/Dao-AILab/flash-attention/blob/main/csrc/flash_attn/src/kernel_traits.h + * Copyright (c) 2025, FourOverSix Team. + ******************************************************************************/ + +#pragma once + +#include "cute/tensor.hpp" + +#include "cutlass/cutlass.h" +#include "cutlass/layout/layout.h" +#include + +using namespace cute; + +template +struct Base_kernel_traits +{ + + static constexpr float E2M1_MAX_VALUE = 6.0f; + static constexpr float E4M3_MIN_VALUE = -448.0f; + static constexpr float E4M3_MAX_VALUE = 448.0f; + static constexpr float E4M3_MAX_FOUROVERSIX = 256.0f; + + using Element = elem_type; + using ScaleFactor = std::conditional_t; + // using ElementXe2m1Packed = std::conditional_t; + using NormConst = float; + static constexpr bool Has_cp_async = true; + + using index_t = int64_t; + + using MMA_Atom_Arch = std::conditional_t< + std::is_same_v, + MMA_Atom, + MMA_Atom>; + + using SmemCopyAtom = Copy_Atom; + using SmemCopyAtomTransposed = Copy_Atom; +}; + +// If Share_Q_K_smem is true, that forces Is_Q_in_regs to be true +template > +struct FP4_quant_kernel_traits : public Base +{ + static constexpr float E2M1_MAX_VALUE = Base::E2M1_MAX_VALUE; + static constexpr float E2M1_MAX_FOUR = 4; + static constexpr float E4M3_MIN_VALUE = Base::E4M3_MIN_VALUE; + static constexpr float E4M3_MAX_VALUE = Base::E4M3_MAX_VALUE; + static constexpr float E4M3_MAX_FOUROVERSIX = Base::E4M3_MAX_FOUROVERSIX; + + using Element = typename Base::Element; + using ScaleFactor = typename Base::ScaleFactor; + using NormConst = typename Base::NormConst; + using index_t = typename Base::index_t; + static constexpr bool Has_cp_async = Base::Has_cp_async; + using SmemCopyAtom = typename Base::SmemCopyAtom; + using SmemCopyAtomTransposed = typename Base::SmemCopyAtomTransposed; + + // The number of threads. + static constexpr int kNWarps = kNWarps_; + static constexpr int kNThreads = kNWarps * 32; + + static constexpr int kBlockM = kBlockM_; + static constexpr int kBlockN = kBlockN_; + static constexpr int kGroupN = Is_nvfp4 ? 16 : 32; // 16 or 32 elements + static_assert(kBlockM % 128 == 0); + static_assert(kBlockN % (kGroupN * 4) == 0); + static_assert(kBlockN % 64 == 0); + static constexpr int kBlockNSmem = 64; // each cache line is 128 bytes, so we need to align to 64 bytes + static constexpr int kBlockNGmem = kBlockN % 128 == 0 ? 128 : 64; + + using TiledMma = TiledMMA< + typename Base::MMA_Atom_Arch, + Layout, _1, _1>>, // 4x1x1 or 8x1x1 thread group + Tile, _16, _16>>; + + // static constexpr int kGroupN = Is_nvfp4 ? 16 : 32; // 16 or 32 elements + // static constexpr int kSwizzleM = Is_nvfp4 ? 4 : 5; // 16 or 32 elements + // static constexpr int kSwizzleS = Is_nvfp4 ? 2 : 1; // 4 or 2 elements + // static constexpr int kSwizzleB = Is_nvfp4 ? 2 : 1; // 2 or 1 bits + + static constexpr int kNumGroupsInRow = kBlockN / kGroupN; + static_assert(kBlockM % kGroupN == 0, "kBlockM must be a multiple of kGroupN if is 2d"); + static constexpr int kNumGroupsInCol = kBlockM; // for 2d scale factor + // static constexpr int kSwizzleM = 3; + // static constexpr int kSwizzleS = 3; + // static constexpr int kSwizzleB = 2; + + // using SmemLayoutAtomX = decltype( + // composition(Swizzle{}, + // // This has to be kBlockNSmem, using kHeadDim gives wrong results for d=128 + // Layout>, + // Stride, _1>>{})); + // using SmemLayoutX = decltype(tile_to_shape( + // SmemLayoutAtomX{}, + // Shape, Int>{})); + using SmemLayoutX = Layout, Int>, Stride, _1>>; + + using SmemLayoutXTransposed = decltype(composition(SmemLayoutX{}, make_layout(Shape, Int>{}, GenRowMajor{}))); + using SmemLayoutXTransposedNoSwizzle = decltype(get_nonswizzle_portion(SmemLayoutXTransposed{})); + + using SmemLayoutSFT = Layout, Int>, + Stride, _1>>; + + static constexpr int kBlockMSF = kBlockM / 128 * 32 * int(kBlockN / (kGroupN * 4)); + static constexpr int kBlockNSF = 16; + using SmemLayoutSF = Layout, Int>, + Stride, _1>>; + + using SmemLayout = SmemLayoutX; + static constexpr int kSmemXSize = size(SmemLayout{}) * sizeof(Element); + static constexpr int kSmemXe2m1Size = kSmemXSize / 4; + static constexpr int kSmemSFTSize = size(SmemLayoutSFT{}) * sizeof(float); + static constexpr int kSmemSFSize = size(SmemLayoutSF{}) * sizeof(ScaleFactor); + static constexpr int kSmemSize = kSmemXSize + kSmemXe2m1Size + kSmemSFTSize + kSmemSFSize; + + static constexpr int kGmemElemsPerLoad = sizeof(cute::uint128_t) / sizeof(Element); + static_assert(kBlockN % kGmemElemsPerLoad == 0, "kBlockN must be a multiple of kGmemElemsPerLoad"); + static constexpr int kGmemThreadsPerRow = kBlockNSmem / kGmemElemsPerLoad; + static_assert(kNThreads % kGmemThreadsPerRow == 0, "kNThreads must be a multiple of kGmemThreadsPerRow"); + using GmemLayoutAtomX = Layout, Int>, + Stride, _1>>; + + using Gmem_copy_atom_x = Copy_Atom, Element>; + + using GmemTiledCopyX = decltype(make_tiled_copy(Gmem_copy_atom_x{}, + GmemLayoutAtomX{}, + Layout>{})); + + using Gmem_copy_atom_sft = Copy_Atom, float>; + + using GmemLayoutAtomSFT = Layout< + Shape<_64, _1>, + Stride<_1, _0>>; + + using GmemTiledCopySFT = decltype(make_tiled_copy(Gmem_copy_atom_sft{}, GmemLayoutAtomSFT{}, Layout>{})); +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/fouroversix/src/fouroversix/csrc/include/static_switch.h b/fouroversix/src/fouroversix/csrc/include/static_switch.h new file mode 100644 index 0000000..38cb2c6 --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/include/static_switch.h @@ -0,0 +1,61 @@ +// Inspired by +// https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h +// and https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Dispatch.h + +// Adapted by Junxian Guo from https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h +// Copyright (c) 2025, FourOverSix Team. + +#pragma once + +/// @param COND - a boolean expression to switch by +/// @param CONST_NAME - a name given for the constexpr bool variable. +/// @param ... - code to execute for true and false +/// +/// Usage: +/// ``` +/// BOOL_SWITCH(flag, BoolConst, [&] { +/// some_function(...); +/// }); +/// ``` + +#define BOOL_SWITCH(COND, CONST_NAME, ...) \ + [&] { \ + if (COND) { \ + constexpr static bool CONST_NAME = true; \ + return __VA_ARGS__(); \ + } else { \ + constexpr static bool CONST_NAME = false; \ + return __VA_ARGS__(); \ + } }() + +#define FP16_SWITCH(COND, ...) \ + [&] { \ + if (COND) { \ + using fp16_type = cutlass::half_t; \ + return __VA_ARGS__(); \ + } else { \ + using fp16_type = cutlass::bfloat16_t; \ + return __VA_ARGS__(); \ + } }() + +#define SELECTION_RULE_SWITCH(SELECTION_RULE, ...) \ + [&] { \ + if (SELECTION_RULE == 0) { \ + constexpr static int kSelectionRule = 0; \ + return __VA_ARGS__(); \ + } else if (SELECTION_RULE == 1) { \ + constexpr static int kSelectionRule = 1; \ + return __VA_ARGS__(); \ + } else if (SELECTION_RULE == 2) { \ + constexpr static int kSelectionRule = 2; \ + return __VA_ARGS__(); \ + } else if (SELECTION_RULE == 3) { \ + constexpr static int kSelectionRule = 3; \ + return __VA_ARGS__(); \ + } else if (SELECTION_RULE == 4) { \ + constexpr static int kSelectionRule = 4; \ + return __VA_ARGS__(); \ + } else { \ + constexpr static int kSelectionRule = 0; \ + return __VA_ARGS__(); \ + } }() diff --git a/fouroversix/src/fouroversix/csrc/include/utils.h b/fouroversix/src/fouroversix/csrc/include/utils.h new file mode 100644 index 0000000..0af1817 --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/include/utils.h @@ -0,0 +1,701 @@ +/****************************************************************************** + * Copyright (c) 2023, Tri Dao. + * Adapted by Junxian Guo and Jack Cook from https://github.com/Dao-AILab/flash-attention/blob/main/csrc/flash_attn/src/utils.h + * Copyright (c) 2025, FourOverSix Team. + ******************************************************************************/ + +#pragma once + +#include +#include +#include + +#include +#include + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 +#include +#endif + +#include + +#include +#include +#include +#include + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace fouroversix +{ + + //////////////////////////////////////////////////////////////////////////////////////////////////// + + template + __forceinline__ __device__ uint32_t relu2(const uint32_t x); + + template <> + __forceinline__ __device__ uint32_t relu2(const uint32_t x) + { + uint32_t res; + const uint32_t zero = 0u; +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + asm volatile("max.f16x2 %0, %1, %2;\n" : "=r"(res) : "r"(x), "r"(zero)); +#else + asm volatile( + "{\n" + "\t .reg .f16x2 sela;\n" + "\t set.gtu.u32.f16x2 sela, %1, %2;\n" + "\t and.b32 %0, sela, %1;\n" + "}\n" : "=r"(res) : "r"(x), "r"(zero)); +#endif + return res; + } + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + template <> + __forceinline__ __device__ uint32_t relu2(const uint32_t x) + { + uint32_t res; + const uint32_t zero = 0u; + asm volatile("max.bf16x2 %0, %1, %2;\n" : "=r"(res) : "r"(x), "r"(zero)); + return res; + } +#endif + + //////////////////////////////////////////////////////////////////////////////////////////////////// + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + + template + __forceinline__ __device__ uint32_t convert_relu2(const float2 x); + + template <> + __forceinline__ __device__ uint32_t convert_relu2(const float2 x) + { + uint32_t res; + const uint32_t a = reinterpret_cast(x.x); + const uint32_t b = reinterpret_cast(x.y); + asm volatile("cvt.rn.relu.f16x2.f32 %0, %1, %2;\n" : "=r"(res) : "r"(b), "r"(a)); + return res; + } + + template <> + __forceinline__ __device__ uint32_t convert_relu2(const float2 x) + { + uint32_t res; + const uint32_t a = reinterpret_cast(x.x); + const uint32_t b = reinterpret_cast(x.y); + asm volatile("cvt.rn.relu.bf16x2.f32 %0, %1, %2;\n" : "=r"(res) : "r"(b), "r"(a)); + return res; + } + +#endif + + //////////////////////////////////////////////////////////////////////////////////////////////////// + + template + struct MaxOp + { + __device__ __forceinline__ T operator()(T const &x, T const &y) { return x > y ? x : y; } + }; + + template <> + struct MaxOp + { + // This is slightly faster + __device__ __forceinline__ float operator()(float const &x, float const &y) { return max(x, y); } + }; + + //////////////////////////////////////////////////////////////////////////////////////////////////// + + template + struct SumOp + { + __device__ __forceinline__ T operator()(T const &x, T const &y) { return x + y; } + }; + + //////////////////////////////////////////////////////////////////////////////////////////////////// + + template + struct Allreduce + { + static_assert(THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4); + template + static __device__ __forceinline__ T run(T x, Operator &op) + { + constexpr int OFFSET = THREADS / 2; + x = op(x, __shfl_xor_sync(uint32_t(-1), x, OFFSET)); + return Allreduce::run(x, op); + } + }; + + //////////////////////////////////////////////////////////////////////////////////////////////////// + + template <> + struct Allreduce<2> + { + template + static __device__ __forceinline__ T run(T x, Operator &op) + { + x = op(x, __shfl_xor_sync(uint32_t(-1), x, 1)); + return x; + } + }; + + //////////////////////////////////////////////////////////////////////////////////////////////////// + + template + __forceinline__ __device__ void gemm(Tensor0 &acc, Tensor1 &tCrA, Tensor2 &tCrB, Tensor3 const &tCsA, + Tensor4 const &tCsB, TiledMma tiled_mma, + TiledCopyA smem_tiled_copy_A, TiledCopyB smem_tiled_copy_B, + ThrCopyA smem_thr_copy_A, ThrCopyB smem_thr_copy_B) + { + CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc)); // MMA_M + CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc)); // MMA_N + CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB)); // MMA_K + Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA); + CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view)); // M + Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB); + CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view)); // N + if (!A_in_regs) + { + cute::copy(smem_tiled_copy_A, tCsA(_, _, _0{}), tCrA_copy_view(_, _, _0{})); + } + if (!B_in_regs) + { + cute::copy(smem_tiled_copy_B, tCsB(_, _, _0{}), tCrB_copy_view(_, _, _0{})); + } +#pragma unroll + for (int i = 0; i < size<2>(tCrA); ++i) + { + if (i < size<2>(tCrA) - 1) + { + if (!A_in_regs) + { + cute::copy(smem_tiled_copy_A, tCsA(_, _, i + 1), tCrA_copy_view(_, _, i + 1)); + } + if (!B_in_regs) + { + cute::copy(smem_tiled_copy_B, tCsB(_, _, i + 1), tCrB_copy_view(_, _, i + 1)); + } + } + cute::gemm(tiled_mma, tCrA(_, _, i), tCrB(_, _, i), acc); + } + } + + //////////////////////////////////////////////////////////////////////////////////////////////////// + + template + __forceinline__ __device__ void gemm_rs(Tensor0 &acc, Tensor1 &tCrA, Tensor2 &tCrB, Tensor3 const &tCsB, + TiledMma tiled_mma, TiledCopy smem_tiled_copy_B, + ThrCopy smem_thr_copy_B) + { + CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc)); // MMA_M + CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc)); // MMA_N + CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB)); // MMA_K + Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB); + CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view)); // N + cute::copy(smem_tiled_copy_B, tCsB(_, _, _0{}), tCrB_copy_view(_, _, _0{})); +#pragma unroll + for (int i = 0; i < size<2>(tCrA); ++i) + { + if (i < size<2>(tCrA) - 1) + { + cute::copy(smem_tiled_copy_B, tCsB(_, _, i + 1), tCrB_copy_view(_, _, i + 1)); + } + cute::gemm(tiled_mma, tCrA(_, _, i), tCrB(_, _, i), acc); + } + } + + //////////////////////////////////////////////////////////////////////////////////////////////////// + + // Convert acc_layout from (MMA=4, MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, MMA_N)) + template + __forceinline__ __device__ auto convert_layout_acc_rowcol(Layout acc_layout) + { + static_assert(decltype(size<0>(acc_layout))::value == 4); + static_assert(decltype(rank(acc_layout))::value == 3); + auto l = logical_divide(acc_layout, Shape<_2>{}); // ((2, 2), MMA_M, MMA_N) + return make_layout(make_layout(get<0, 1>(l), get<1>(l)), make_layout(get<0, 0>(l), get<2>(l))); + }; + + //////////////////////////////////////////////////////////////////////////////////////////////////// + + // Convert acc_layout from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) + // if using m16n8k16, or to (4, MMA_M, MMA_N) if using m16n8k8. + template + __forceinline__ __device__ auto convert_layout_acc_Aregs(Layout acc_layout) + { + using X = Underscore; + static_assert(decltype(size<0>(acc_layout))::value == 4); + static_assert(decltype(rank(acc_layout))::value == 3); + constexpr int mma_shape_K = get<2>(typename MMA_traits::Shape_MNK{}); + static_assert(mma_shape_K == 8 || mma_shape_K == 16); + if constexpr (mma_shape_K == 8) + { + return acc_layout; + } + else + { + auto l = logical_divide(acc_layout, Shape{}); // (4, MMA_M, (2, MMA_N / 2))) + return make_layout(make_layout(get<0>(l), get<2, 0>(l)), get<1>(l), get<2, 1>(l)); + } + }; + + //////////////////////////////////////////////////////////////////////////////////////////////////// + + // Convert acc_layout from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) + template + __forceinline__ __device__ auto convert_layout_acc_dropout(Layout acc_layout) + { + using X = Underscore; + static_assert(decltype(size<0>(acc_layout))::value == 4); + static_assert(decltype(rank(acc_layout))::value == 3); + auto l = logical_divide(acc_layout, Shape{}); // (4, MMA_M, (2, MMA_N / 2))) + return make_layout(make_layout(get<0>(l), get<2, 0>(l)), get<1>(l), get<2, 1>(l)); + }; + + //////////////////////////////////////////////////////////////////////////////////////////////////// + + template + __forceinline__ __device__ auto convert_type(Tensor const &tensor) + { + using From_type = typename Engine::value_type; + constexpr int numel = decltype(size(tensor))::value; + cutlass::NumericArrayConverter convert_op; + // HACK: this requires tensor to be "contiguous" + auto frag = convert_op(*reinterpret_cast *>(tensor.data())); + return make_tensor(make_rmem_ptr(&frag), tensor.layout()); + } + + //////////////////////////////////////////////////////////////////////////////////////////////////// + + // Blocks until all but N previous cp.async.commit_group operations have committed. + // This differs from cute::cp_async_wait in that when N = 0 we don't call cp.async.wait_all + // (which is equivalent to commit_group then wait_group 0). + // Instead we just call cp.async.wait_group 0, which is slightly faster. + // https://github.com/NVIDIA/cutlass/blob/master/include/cute/arch/copy_sm80.hpp#L113 + template + CUTE_HOST_DEVICE void cp_async_wait() + { +#if defined(CUTE_ARCH_CP_ASYNC_SM80_ENABLED) + asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); +#endif + } + + //////////////////////////////////////////////////////////////////////////////////////////////////// + + template + __forceinline__ __device__ void copy(TiledCopy tiled_copy, Tensor const &S, + Tensor &D, Tensor const &identity_MN, + Tensor const &predicate_K, const int max_MN = 0) + { + CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{}); + CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{}); + CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); // MMA + CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); // MMA_M + CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); // MMA_K + // There's no case where !Clear_OOB_K && Clear_OOB_MN + static_assert(!(Clear_OOB_MN && !Clear_OOB_K)); +#pragma unroll + for (int m = 0; m < size<1>(S); ++m) + { + if (Is_even_MN || get<0>(identity_MN(0, m, 0)) < max_MN) + { +#pragma unroll + for (int k = 0; k < size<2>(S); ++k) + { + if (Is_even_K || predicate_K(k)) + { + cute::copy(tiled_copy, S(_, m, k), D(_, m, k)); + } + else if (Clear_OOB_K) + { + cute::clear(D(_, m, k)); + } + } + } + else if (Clear_OOB_MN) + { + cute::clear(D(_, m, _)); + } + } + } + + static __device__ __forceinline__ float atomicMaxFloat(float *addr, float value) + { + // source: https://stackoverflow.com/a/51549250 + return (value >= 0) + ? __int_as_float(atomicMax((int *)addr, __float_as_int(value))) + : __uint_as_float(atomicMin((unsigned int *)addr, __float_as_uint(value))); + } + + //////////////////////////////////////////////////////////////////////////////////////////////////// + + template + struct Fp4ArrayQuant + { + using InputType = cutlass::Array; + using OutputType = cutlass::Array; + using ScaleFactorType = float; + using ErrorType = float; + + __device__ __forceinline__ + OutputType + convert(InputType const &x, + const float amax, + const ScaleFactorType sf, + const uint32_t rbits, + // Usage depends on Is_rtn + ErrorType *err /*nullable*/) + { + InputType x_scaled; + constexpr float E2M1_MAX_VALUE = 6.0f; + constexpr float E2M1_MAX_FOUR = 4.0f; + constexpr float E4M3_MAX_VALUE = 448.0f; + constexpr float E4M3_MAX_FOUROVERSIX = 256.0f; + + constexpr float e2m1_limit = kAdaptiveBlockScalingRuleType == AdaptiveBlockScalingRuleType::STATIC_4 ? E2M1_MAX_FOUR : E2M1_MAX_VALUE; + constexpr float e4m3_limit = (kAdaptiveBlockScalingRuleType == AdaptiveBlockScalingRuleType::STATIC_6 || kAdaptiveBlockScalingRuleType == AdaptiveBlockScalingRuleType::STATIC_4) + ? E4M3_MAX_VALUE + : E4M3_MAX_FOUROVERSIX; + const float encode_scale = e4m3_limit * e2m1_limit / amax; + const float decode_scale = 1.0 / encode_scale; + const float block_scale_inv = fminf(1.0f / (decode_scale * sf), std::numeric_limits::max()); + +#pragma unroll + for (int i = 0; i < 8; ++i) + { + x_scaled[i] = x[i] * block_scale_inv; + } + + unsigned out; + + if constexpr (Is_rtn) + { + if constexpr (Is_4o6) + { + unsigned out_dequant_1; + unsigned out_dequant_2; + unsigned out_dequant_3; + unsigned out_dequant_4; + + asm volatile( + "{\n" + ".reg .b8 byte0, byte1, byte2, byte3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %6, %5;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte1, %8, %7;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte2, %10, %9;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte3, %12, %11;\n" + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" + "cvt.rn.f16x2.e2m1x2 %1, byte0;\n" + "cvt.rn.f16x2.e2m1x2 %2, byte1;\n" + "cvt.rn.f16x2.e2m1x2 %3, byte2;\n" + "cvt.rn.f16x2.e2m1x2 %4, byte3;\n" + "}" + : "=r"(out), "=r"(out_dequant_1), "=r"(out_dequant_2), "=r"(out_dequant_3), "=r"(out_dequant_4) : "f"(x_scaled[0]), "f"(x_scaled[1]), "f"(x_scaled[2]), "f"(x_scaled[3]), + "f"(x_scaled[4]), "f"(x_scaled[5]), "f"(x_scaled[6]), "f"(x_scaled[7])); + + unsigned short out_dequant_1_hi = (out_dequant_1 >> 16) & 0xFFFF; + unsigned short out_dequant_1_lo = out_dequant_1 & 0xFFFF; + unsigned short out_dequant_2_hi = (out_dequant_2 >> 16) & 0xFFFF; + unsigned short out_dequant_2_lo = out_dequant_2 & 0xFFFF; + unsigned short out_dequant_3_hi = (out_dequant_3 >> 16) & 0xFFFF; + unsigned short out_dequant_3_lo = out_dequant_3 & 0xFFFF; + unsigned short out_dequant_4_hi = (out_dequant_4 >> 16) & 0xFFFF; + unsigned short out_dequant_4_lo = out_dequant_4 & 0xFFFF; + + float val0 = __half2float(__ushort_as_half(out_dequant_1_lo)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX); + float val1 = __half2float(__ushort_as_half(out_dequant_1_hi)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX); + float val2 = __half2float(__ushort_as_half(out_dequant_2_lo)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX); + float val3 = __half2float(__ushort_as_half(out_dequant_2_hi)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX); + float val4 = __half2float(__ushort_as_half(out_dequant_3_lo)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX); + float val5 = __half2float(__ushort_as_half(out_dequant_3_hi)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX); + float val6 = __half2float(__ushort_as_half(out_dequant_4_lo)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX); + float val7 = __half2float(__ushort_as_half(out_dequant_4_hi)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX); + + if constexpr (kAdaptiveBlockScalingRuleType == AdaptiveBlockScalingRuleType::MAE_4o6) + { + *err += std::abs(val0 - x[0]); + *err += std::abs(val1 - x[1]); + *err += std::abs(val2 - x[2]); + *err += std::abs(val3 - x[3]); + *err += std::abs(val4 - x[4]); + *err += std::abs(val5 - x[5]); + *err += std::abs(val6 - x[6]); + *err += std::abs(val7 - x[7]); + } + else if constexpr (kAdaptiveBlockScalingRuleType == AdaptiveBlockScalingRuleType::MSE_4o6) + { + *err += (val0 - x[0]) * (val0 - x[0]); + *err += (val1 - x[1]) * (val1 - x[1]); + *err += (val2 - x[2]) * (val2 - x[2]); + *err += (val3 - x[3]) * (val3 - x[3]); + *err += (val4 - x[4]) * (val4 - x[4]); + *err += (val5 - x[5]) * (val5 - x[5]); + *err += (val6 - x[6]) * (val6 - x[6]); + *err += (val7 - x[7]) * (val7 - x[7]); + } + else if constexpr (kAdaptiveBlockScalingRuleType == AdaptiveBlockScalingRuleType::ABS_MAX_4o6) + { + float val0_err = std::abs(val0 - x[0]); + if (val0_err > *err) + *err = val0_err; + float val1_err = std::abs(val1 - x[1]); + if (val1_err > *err) + *err = val1_err; + float val2_err = std::abs(val2 - x[2]); + if (val2_err > *err) + *err = val2_err; + float val3_err = std::abs(val3 - x[3]); + if (val3_err > *err) + *err = val3_err; + float val4_err = std::abs(val4 - x[4]); + if (val4_err > *err) + *err = val4_err; + float val5_err = std::abs(val5 - x[5]); + if (val5_err > *err) + *err = val5_err; + float val6_err = std::abs(val6 - x[6]); + if (val6_err > *err) + *err = val6_err; + float val7_err = std::abs(val7 - x[7]); + if (val7_err > *err) + *err = val7_err; + } + else + { + printf("in Fp4ArrayQuant::convert, kAdaptiveBlockScalingRuleType = %d, not supported\n", kAdaptiveBlockScalingRuleType); + assert(false); + } + + return reinterpret_cast(out); + } + else + { + asm volatile( + "{\n" + ".reg .b8 byte0, byte1, byte2, byte3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" + "}" + : "=r"(out) : "f"(x_scaled[0]), "f"(x_scaled[1]), "f"(x_scaled[2]), "f"(x_scaled[3]), + "f"(x_scaled[4]), "f"(x_scaled[5]), "f"(x_scaled[6]), "f"(x_scaled[7])); + return reinterpret_cast(out); + } + } + else + { + if constexpr (Is_4o6) + { + unsigned out_dequant_1; + unsigned out_dequant_2; + unsigned out_dequant_3; + unsigned out_dequant_4; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030) + asm volatile( + "{\n" + ".reg .b16 tmp0, tmp1;\n" + ".reg .b8 byte0, byte1;\n" + "cvt.rs.satfinite.e2m1x4.f32 tmp0, {%8, %7, %6, %5}, %13;\n" + "mov.b16 {byte1, byte0}, tmp0;\n" + "cvt.rn.f16x2.e2m1x2 %1, byte0;\n" + "cvt.rn.f16x2.e2m1x2 %2, byte1;\n" + "cvt.rs.satfinite.e2m1x4.f32 tmp1, {%12, %11, %10, %9}, %14;\n" + "mov.b16 {byte1, byte0}, tmp1;\n" + "cvt.rn.f16x2.e2m1x2 %3, byte0;\n" + "cvt.rn.f16x2.e2m1x2 %4, byte1;\n" + "mov.b32 %0, {tmp0, tmp1};\n" + "}" + : "=r"(out), "=r"(out_dequant_1), "=r"(out_dequant_2), "=r"(out_dequant_3), "=r"(out_dequant_4) : "f"(x_scaled[0]), "f"(x_scaled[1]), "f"(x_scaled[2]), "f"(x_scaled[3]), "f"(x_scaled[4]), "f"(x_scaled[5]), "f"(x_scaled[6]), "f"(x_scaled[7]), "r"(rbits), "r"(rbits)); +#endif + + unsigned short out_dequant_1_hi = (out_dequant_1 >> 16) & 0xFFFF; + unsigned short out_dequant_1_lo = out_dequant_1 & 0xFFFF; + unsigned short out_dequant_2_hi = (out_dequant_2 >> 16) & 0xFFFF; + unsigned short out_dequant_2_lo = out_dequant_2 & 0xFFFF; + unsigned short out_dequant_3_hi = (out_dequant_3 >> 16) & 0xFFFF; + unsigned short out_dequant_3_lo = out_dequant_3 & 0xFFFF; + unsigned short out_dequant_4_hi = (out_dequant_4 >> 16) & 0xFFFF; + unsigned short out_dequant_4_lo = out_dequant_4 & 0xFFFF; + + float val0 = __half2float(__ushort_as_half(out_dequant_1_lo)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX); + float val1 = __half2float(__ushort_as_half(out_dequant_1_hi)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX); + float val2 = __half2float(__ushort_as_half(out_dequant_2_lo)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX); + float val3 = __half2float(__ushort_as_half(out_dequant_2_hi)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX); + float val4 = __half2float(__ushort_as_half(out_dequant_3_lo)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX); + float val5 = __half2float(__ushort_as_half(out_dequant_3_hi)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX); + float val6 = __half2float(__ushort_as_half(out_dequant_4_lo)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX); + float val7 = __half2float(__ushort_as_half(out_dequant_4_hi)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX); + + if constexpr (kAdaptiveBlockScalingRuleType == AdaptiveBlockScalingRuleType::MAE_4o6) + { + *err += std::abs(val0 - x[0]); + *err += std::abs(val1 - x[1]); + *err += std::abs(val2 - x[2]); + *err += std::abs(val3 - x[3]); + *err += std::abs(val4 - x[4]); + *err += std::abs(val5 - x[5]); + *err += std::abs(val6 - x[6]); + *err += std::abs(val7 - x[7]); + } + else if constexpr (kAdaptiveBlockScalingRuleType == AdaptiveBlockScalingRuleType::MSE_4o6) + { + *err += (val0 - x[0]) * (val0 - x[0]); + *err += (val1 - x[1]) * (val1 - x[1]); + *err += (val2 - x[2]) * (val2 - x[2]); + *err += (val3 - x[3]) * (val3 - x[3]); + *err += (val4 - x[4]) * (val4 - x[4]); + *err += (val5 - x[5]) * (val5 - x[5]); + *err += (val6 - x[6]) * (val6 - x[6]); + *err += (val7 - x[7]) * (val7 - x[7]); + } + else if constexpr (kAdaptiveBlockScalingRuleType == AdaptiveBlockScalingRuleType::ABS_MAX_4o6) + { + float val0_err = std::abs(val0 - x[0]); + if (val0_err > *err) + *err = val0_err; + float val1_err = std::abs(val1 - x[1]); + if (val1_err > *err) + *err = val1_err; + float val2_err = std::abs(val2 - x[2]); + if (val2_err > *err) + *err = val2_err; + float val3_err = std::abs(val3 - x[3]); + if (val3_err > *err) + *err = val3_err; + float val4_err = std::abs(val4 - x[4]); + if (val4_err > *err) + *err = val4_err; + float val5_err = std::abs(val5 - x[5]); + if (val5_err > *err) + *err = val5_err; + float val6_err = std::abs(val6 - x[6]); + if (val6_err > *err) + *err = val6_err; + float val7_err = std::abs(val7 - x[7]); + if (val7_err > *err) + *err = val7_err; + } + else + { + printf("in Fp4ArrayQuant::convert, kAdaptiveBlockScalingRuleType = %d, not supported\n", kAdaptiveBlockScalingRuleType); + assert(false); + } + + return reinterpret_cast(out); + } + else + { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030) + asm volatile( + "{\n" + ".reg .b16 tmp0, tmp1;\n" + "cvt.rs.satfinite.e2m1x4.f32 tmp0, {%4, %3, %2, %1}, %9;\n" + "cvt.rs.satfinite.e2m1x4.f32 tmp1, {%8, %7, %6, %5}, %10;\n" + "mov.b32 %0, {tmp0, tmp1};\n" + "}" + : "=r"(out) : "f"(x_scaled[0]), "f"(x_scaled[1]), "f"(x_scaled[2]), "f"(x_scaled[3]), + "f"(x_scaled[4]), "f"(x_scaled[5]), "f"(x_scaled[6]), "f"(x_scaled[7]), + "r"(rbits), "r"(rbits)); +#endif + return reinterpret_cast(out); + } + } + } + }; + + template + __forceinline__ __device__ float fp4_conversion(Tensor const &tensor, const float amax, float *sf_, OutputType *res, const uint32_t rbits) + { + constexpr int numel = decltype(size(tensor))::value; + static_assert((numel == 16 && Is_nvfp4) || numel == 32); + static_assert(std::is_same_v>); + + constexpr int loop_size = 8; + constexpr int num_loops = numel / loop_size; + + using InputType = cutlass::Array; + + Fp4ArrayQuant fp4_array_quant; + + if constexpr (Is_4o6) + { + float err[2] = {0.0f, 0.0f}; + OutputType res_4[num_loops]; + OutputType res_6[num_loops]; + float final_err[2] = {0.0f, 0.0f}; + +#pragma unroll + for (int i = 0; i < num_loops; ++i) + { + InputType x; +#pragma unroll + for (int j = 0; j < loop_size; ++j) + { + x[j] = static_cast(tensor(i * loop_size + j)); + } + res_4[i] = fp4_array_quant.convert(x, amax, sf_[0], rbits, &err[0]); + res_6[i] = fp4_array_quant.convert(x, amax, sf_[1], rbits, &err[1]); + } + + if (Is_2d){ + // For 2D tensors we want to pick the same format for the entire tensor, to keep it simple for downstream processing. + // So we pick the format with smaller total error across the entire tensor. + + // If the method is MAE or MSE, we need to sum the error across the entire tensor. If the method is ABS_MAX, we need to take the max error across the entire tensor. + using RedOp = std::conditional_t, SumOp>; + RedOp op; + final_err[0] = Allreduce::run(err[0], op); + final_err[1] = Allreduce::run(err[1], op); + } else { + final_err[0] = err[0]; + final_err[1] = err[1]; + } + + // pick_first = true means choose 4, false means choose 6 + bool const pick_first = final_err[0] < final_err[1]; +#pragma unroll + for (int i = 0; i < num_loops; ++i) + { + res[i] = pick_first ? res_4[i] : res_6[i]; + } + return sf_[!pick_first]; + } + else + { +#pragma unroll + for (int i = 0; i < num_loops; ++i) + { + InputType x; +#pragma unroll + for (int j = 0; j < loop_size; ++j) + { + x[j] = static_cast(tensor(i * loop_size + j)); + } + res[i] = fp4_array_quant.convert(x, amax, sf_[0], rbits, nullptr); + } + + return sf_[0]; + } + } + + __device__ __forceinline__ float clamp(float value, float min_value, float max_value) + { + return max(min(value, max_value), min_value); + } + + //////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace fouroversix diff --git a/fouroversix/src/fouroversix/csrc/quantize/fp4_quant.cu b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant.cu new file mode 100644 index 0000000..2eb1682 --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant.cu @@ -0,0 +1,245 @@ +#include +#include + +#include +#include +#include +#include +#include // For at::Generator and at::PhiloxCudaState + +#include + +#include "hardware_info.h" +#include "fp4_quant.h" +#include "static_switch.h" + +#define CHECK_DEVICE(x) TORCH_CHECK(x.is_cuda(), #x " must be on CUDA") +#define CHECK_SHAPE(x, ...) TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")") +#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") + +namespace fouroversix +{ + void set_params_fp4_quant( + FP4_quant_params ¶ms, + /*-------------- tensors ---------------*/ + const at::Tensor x, + at::Tensor x_rht, + at::Tensor x_e2m1, + at::Tensor x_sf, + at::Tensor x_sft, + at::Tensor amax, + const int M, + const int N, + const int M_rounded, + const int N_rounded, + const int M_sf, + const int N_sf, + const bool is_nvfp4, + const bool is_rtn, + const bool is_rht, + const bool is_2d, + // const bool is_4o6, + const bool is_transpose, + const int selection_rule, + const int rbits) + { + + // Reset the parameters + params = {}; + + params.is_bf16 = x.dtype() == torch::kBFloat16; + + /**************** Pointers & strides ****************/ + params.x_ptr = x.data_ptr(); + params.x_rht_ptr = x_rht.data_ptr(); + params.x_e2m1_ptr = x_e2m1.data_ptr(); + params.x_sf_ptr = x_sf.data_ptr(); + params.x_sft_ptr = x_sft.data_ptr(); + params.amax_ptr = amax.data_ptr(); + + // Element-based strides (not bytes) + params.x_row_stride = x.stride(0); + params.x_col_stride = x.stride(1); + params.x_rht_row_stride = x_rht.stride(0); + params.x_rht_col_stride = x_rht.stride(1); + params.x_e2m1_row_stride = x_e2m1.stride(0); + params.x_e2m1_col_stride = x_e2m1.stride(1); + params.x_sf_row_stride = x_sf.stride(0); + params.x_sf_col_stride = x_sf.stride(1); + params.x_sft_row_stride = x_sft.stride(0); + params.x_sft_col_stride = x_sft.stride(1); + + // Set the dimensions + params.M = M; + params.N = N; + params.M_rounded = M_rounded; + params.N_rounded = N_rounded; + params.M_sf = M_sf; + params.N_sf = N_sf; + // Set FP4-specific parameters + params.is_nvfp4 = is_nvfp4; + params.is_rtn = is_rtn; + params.is_rht = is_rht; + params.is_2d = is_2d; + // params.is_4o6 = is_4o6; + params.is_transpose = is_transpose; + params.selection_rule = selection_rule; + params.rbits = rbits; + } + + void run_fp4_quant(FP4_quant_params ¶ms, cudaStream_t stream) + { + FP16_SWITCH(!params.is_bf16, [&] + { BOOL_SWITCH(params.is_nvfp4, Is_nvfp4, [&] + { BOOL_SWITCH(params.is_rht, Is_rht, [&] + { BOOL_SWITCH(params.is_transpose, Is_transpose, [&] + { run_fp4_quant_(params, stream); }); }); }); }); + } + + std::tuple quantize_to_fp4( + const at::Tensor &x, + const bool is_nvfp4, + const bool is_rtn, + const bool is_rht, + // const bool is_4o6, + const bool is_2d, + const bool is_transpose, + const int64_t selection_rule, + const int64_t rbits) + { + + /******* + * selection_rule: + * 0: static_6 + * 1: static_4 + * 2: 4o6_l1_norm + * 3: 4o6_mse + * 4: 4o6_abs_max + */ + TORCH_CHECK(selection_rule >= 0 && selection_rule <= 4, "Invalid selection_rule: " + std::to_string(selection_rule)); + // const int is_4o6 = selection_rule == 2 || selection_rule == 3; + + /********************** + * 1. Sanity checks * + *********************/ + at::cuda::CUDAGuard device_guard{x.device()}; + + // Hardware capability + { + auto [cc_major, _] = get_compute_capability(get_current_device()); + TORCH_CHECK(cc_major >= 10, "FP4Quant only supports Blackwell GPUs or newer."); + } + + // Dtype / device checks + auto dtype = x.dtype(); + TORCH_CHECK(dtype == torch::kFloat16 || dtype == torch::kBFloat16, + "FP4Quant only supports fp16 and bf16 data types"); + // TORCH_CHECK(km.dtype() == dtype, "q and km must have the same dtype"); + CHECK_DEVICE(x); + // Layout / contiguity checks + TORCH_CHECK(x.stride(-1) == 1, "x must be contiguous on the last dim"); + + /********************** + * 2. Dimension logic * + *********************/ + + const int M = is_transpose ? x.size(1) : x.size(0); + const int N = is_transpose ? x.size(0) : x.size(1); + + const int n_round = is_nvfp4 ? 64 : 128; + TORCH_CHECK(N % n_round == 0, "N must be multiple of " + std::to_string(n_round)); + + /********************** + * 3. Derived sizes * + *********************/ + auto round_up = [](int x, int m) + { return (x + m - 1) / m * m; }; + + const int M_rounded = round_up(M, 128); + const int N_rounded = round_up(N, n_round); + + // const int max_seqlen_km = ceil_div(max_seqlen_k, moba_chunk_size); + // const int head_size_rounded = round_up(head_size, head_size <= 128 ? 32 : 64); + // const int seqlen_q_rounded = round_up(max_seqlen_q, 128); + // const int seqlen_km_rounded = round_up(max_seqlen_km, 128); + // const int moba_topk_rounded = round_up(moba_topk, 16); + + /********************** + * 4. Intermediate buffers * + *********************/ + at::Tensor x_rht; + if (is_rht) + { + x_rht = torch::zeros({M_rounded, N_rounded}, x.options()); + } + else + { + x_rht = torch::zeros({0, 0}, x.options()); + } + + /********************** + * 5. Output buffers * + *********************/ + at::Tensor x_e2m1 = torch::zeros({M_rounded, int(N_rounded / 2)}, x.options().dtype(torch::kUInt8)); + at::Tensor x_sf, x_sft; + int M_sf, N_sf; + if (is_nvfp4) + { + M_sf = int(M_rounded / 128 * 32) * int(N_rounded / 64); + N_sf = 16; + // N_sf = int(N_rounded / 16 * 4); + x_sf = torch::zeros({M_sf, N_sf}, x.options().dtype(torch::kFloat8_e4m3fn)); + x_sft = torch::zeros({M_rounded, int(N_rounded / 16)}, x.options().dtype(torch::kFloat32)); + } + else + { + M_sf = int(M_rounded / 128 * 32) * int(N_rounded / 128); + N_sf = 16; + x_sf = torch::zeros({M_sf, N_sf}, x.options().dtype(torch::kUInt8)); + x_sft = torch::zeros({M_rounded, int(N_rounded / 32)}, x.options().dtype(torch::kFloat32)); + } + at::Tensor amax = torch::zeros({1}, x.options().dtype(torch::kFloat32)); + + /********************** + * 5. Param struct * + *********************/ + FP4_quant_params params; + // const at::Tensor x, + // at::Tensor x_rht, + // at::Tensor x_e2m1, + // at::Tensor x_sf, + // at::Tensor amax, + // const int M, + // const int N, + // const bool is_nvfp4, + // const bool is_rtn, + // const bool is_4o6, + // const bool is_2d, + // const bool is_transpose + set_params_fp4_quant( + params, + /*-------------- tensors ---------------*/ + x, x_rht, x_e2m1, x_sf, x_sft, amax, M, N, M_rounded, N_rounded, M_sf, N_sf, + is_nvfp4, is_rtn, is_rht, is_2d, is_transpose, selection_rule, rbits); + + /********************** + * 6. Kernel launch * + *********************/ + + if (M > 0) + { + run_fp4_quant(params, at::cuda::getCurrentCUDAStream().stream()); + } + else + { + amax.fill_(0); + } + + return std::make_tuple(x_e2m1, x_sf.flatten(), amax); + } + + TORCH_LIBRARY_IMPL(fouroversix, CUDA, m) + { + m.impl("quantize_to_fp4", &quantize_to_fp4); + } +} \ No newline at end of file diff --git a/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_mxfp4_rht_sm100.cu b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_mxfp4_rht_sm100.cu new file mode 100644 index 0000000..fc33a83 --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_mxfp4_rht_sm100.cu @@ -0,0 +1,11 @@ +// Splitting the different transpose modes to different files to speed up compilation. +// This file is auto-generated. See "generate_kernels.py" +#include "fp4_quant_launch_template.h" +namespace fouroversix { + +template<> +void run_fp4_quant_(FP4_quant_params ¶ms, cudaStream_t stream) { + run_mxfp4_quant_rht(params, stream); +} + +} // namespace fouroversix \ No newline at end of file diff --git a/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_mxfp4_rht_trans_sm100.cu b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_mxfp4_rht_trans_sm100.cu new file mode 100644 index 0000000..8fdc539 --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_mxfp4_rht_trans_sm100.cu @@ -0,0 +1,11 @@ +// Splitting the different transpose modes to different files to speed up compilation. +// This file is auto-generated. See "generate_kernels.py" +#include "fp4_quant_launch_template.h" +namespace fouroversix { + +template<> +void run_fp4_quant_(FP4_quant_params ¶ms, cudaStream_t stream) { + run_mxfp4_quant_rht(params, stream); +} + +} // namespace fouroversix \ No newline at end of file diff --git a/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_mxfp4_sm100.cu b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_mxfp4_sm100.cu new file mode 100644 index 0000000..38022fb --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_mxfp4_sm100.cu @@ -0,0 +1,11 @@ +// Splitting the different transpose modes to different files to speed up compilation. +// This file is auto-generated. See "generate_kernels.py" +#include "fp4_quant_launch_template.h" +namespace fouroversix { + +template<> +void run_fp4_quant_(FP4_quant_params ¶ms, cudaStream_t stream) { + run_mxfp4_quant(params, stream); +} + +} // namespace fouroversix \ No newline at end of file diff --git a/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_mxfp4_trans_sm100.cu b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_mxfp4_trans_sm100.cu new file mode 100644 index 0000000..6886f81 --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_mxfp4_trans_sm100.cu @@ -0,0 +1,11 @@ +// Splitting the different transpose modes to different files to speed up compilation. +// This file is auto-generated. See "generate_kernels.py" +#include "fp4_quant_launch_template.h" +namespace fouroversix { + +template<> +void run_fp4_quant_(FP4_quant_params ¶ms, cudaStream_t stream) { + run_mxfp4_quant(params, stream); +} + +} // namespace fouroversix \ No newline at end of file diff --git a/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_nvfp4_rht_sm100.cu b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_nvfp4_rht_sm100.cu new file mode 100644 index 0000000..884bae7 --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_nvfp4_rht_sm100.cu @@ -0,0 +1,11 @@ +// Splitting the different transpose modes to different files to speed up compilation. +// This file is auto-generated. See "generate_kernels.py" +#include "fp4_quant_launch_template.h" +namespace fouroversix { + +template<> +void run_fp4_quant_(FP4_quant_params ¶ms, cudaStream_t stream) { + run_nvfp4_quant_rht(params, stream); +} + +} // namespace fouroversix \ No newline at end of file diff --git a/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_nvfp4_rht_trans_sm100.cu b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_nvfp4_rht_trans_sm100.cu new file mode 100644 index 0000000..4345139 --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_nvfp4_rht_trans_sm100.cu @@ -0,0 +1,11 @@ +// Splitting the different transpose modes to different files to speed up compilation. +// This file is auto-generated. See "generate_kernels.py" +#include "fp4_quant_launch_template.h" +namespace fouroversix { + +template<> +void run_fp4_quant_(FP4_quant_params ¶ms, cudaStream_t stream) { + run_nvfp4_quant_rht(params, stream); +} + +} // namespace fouroversix \ No newline at end of file diff --git a/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_nvfp4_sm100.cu b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_nvfp4_sm100.cu new file mode 100644 index 0000000..ed5bc5b --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_nvfp4_sm100.cu @@ -0,0 +1,11 @@ +// Splitting the different transpose modes to different files to speed up compilation. +// This file is auto-generated. See "generate_kernels.py" +#include "fp4_quant_launch_template.h" +namespace fouroversix { + +template<> +void run_fp4_quant_(FP4_quant_params ¶ms, cudaStream_t stream) { + run_nvfp4_quant(params, stream); +} + +} // namespace fouroversix \ No newline at end of file diff --git a/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_nvfp4_trans_sm100.cu b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_nvfp4_trans_sm100.cu new file mode 100644 index 0000000..3370703 --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_nvfp4_trans_sm100.cu @@ -0,0 +1,11 @@ +// Splitting the different transpose modes to different files to speed up compilation. +// This file is auto-generated. See "generate_kernels.py" +#include "fp4_quant_launch_template.h" +namespace fouroversix { + +template<> +void run_fp4_quant_(FP4_quant_params ¶ms, cudaStream_t stream) { + run_nvfp4_quant(params, stream); +} + +} // namespace fouroversix \ No newline at end of file diff --git a/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_fp16_mxfp4_rht_sm100.cu b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_fp16_mxfp4_rht_sm100.cu new file mode 100644 index 0000000..45a12dd --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_fp16_mxfp4_rht_sm100.cu @@ -0,0 +1,11 @@ +// Splitting the different transpose modes to different files to speed up compilation. +// This file is auto-generated. See "generate_kernels.py" +#include "fp4_quant_launch_template.h" +namespace fouroversix { + +template<> +void run_fp4_quant_(FP4_quant_params ¶ms, cudaStream_t stream) { + run_mxfp4_quant_rht(params, stream); +} + +} // namespace fouroversix \ No newline at end of file diff --git a/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_fp16_mxfp4_rht_trans_sm100.cu b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_fp16_mxfp4_rht_trans_sm100.cu new file mode 100644 index 0000000..463bbea --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_fp16_mxfp4_rht_trans_sm100.cu @@ -0,0 +1,11 @@ +// Splitting the different transpose modes to different files to speed up compilation. +// This file is auto-generated. See "generate_kernels.py" +#include "fp4_quant_launch_template.h" +namespace fouroversix { + +template<> +void run_fp4_quant_(FP4_quant_params ¶ms, cudaStream_t stream) { + run_mxfp4_quant_rht(params, stream); +} + +} // namespace fouroversix \ No newline at end of file diff --git a/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_fp16_mxfp4_sm100.cu b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_fp16_mxfp4_sm100.cu new file mode 100644 index 0000000..d972e32 --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_fp16_mxfp4_sm100.cu @@ -0,0 +1,11 @@ +// Splitting the different transpose modes to different files to speed up compilation. +// This file is auto-generated. See "generate_kernels.py" +#include "fp4_quant_launch_template.h" +namespace fouroversix { + +template<> +void run_fp4_quant_(FP4_quant_params ¶ms, cudaStream_t stream) { + run_mxfp4_quant(params, stream); +} + +} // namespace fouroversix \ No newline at end of file diff --git a/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_fp16_mxfp4_trans_sm100.cu b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_fp16_mxfp4_trans_sm100.cu new file mode 100644 index 0000000..9919a90 --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_fp16_mxfp4_trans_sm100.cu @@ -0,0 +1,11 @@ +// Splitting the different transpose modes to different files to speed up compilation. +// This file is auto-generated. See "generate_kernels.py" +#include "fp4_quant_launch_template.h" +namespace fouroversix { + +template<> +void run_fp4_quant_(FP4_quant_params ¶ms, cudaStream_t stream) { + run_mxfp4_quant(params, stream); +} + +} // namespace fouroversix \ No newline at end of file diff --git a/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_fp16_nvfp4_rht_sm100.cu b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_fp16_nvfp4_rht_sm100.cu new file mode 100644 index 0000000..2375e20 --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_fp16_nvfp4_rht_sm100.cu @@ -0,0 +1,11 @@ +// Splitting the different transpose modes to different files to speed up compilation. +// This file is auto-generated. See "generate_kernels.py" +#include "fp4_quant_launch_template.h" +namespace fouroversix { + +template<> +void run_fp4_quant_(FP4_quant_params ¶ms, cudaStream_t stream) { + run_nvfp4_quant_rht(params, stream); +} + +} // namespace fouroversix \ No newline at end of file diff --git a/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_fp16_nvfp4_rht_trans_sm100.cu b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_fp16_nvfp4_rht_trans_sm100.cu new file mode 100644 index 0000000..9d85f1e --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_fp16_nvfp4_rht_trans_sm100.cu @@ -0,0 +1,11 @@ +// Splitting the different transpose modes to different files to speed up compilation. +// This file is auto-generated. See "generate_kernels.py" +#include "fp4_quant_launch_template.h" +namespace fouroversix { + +template<> +void run_fp4_quant_(FP4_quant_params ¶ms, cudaStream_t stream) { + run_nvfp4_quant_rht(params, stream); +} + +} // namespace fouroversix \ No newline at end of file diff --git a/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_fp16_nvfp4_sm100.cu b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_fp16_nvfp4_sm100.cu new file mode 100644 index 0000000..b6b2cbf --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_fp16_nvfp4_sm100.cu @@ -0,0 +1,11 @@ +// Splitting the different transpose modes to different files to speed up compilation. +// This file is auto-generated. See "generate_kernels.py" +#include "fp4_quant_launch_template.h" +namespace fouroversix { + +template<> +void run_fp4_quant_(FP4_quant_params ¶ms, cudaStream_t stream) { + run_nvfp4_quant(params, stream); +} + +} // namespace fouroversix \ No newline at end of file diff --git a/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_fp16_nvfp4_trans_sm100.cu b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_fp16_nvfp4_trans_sm100.cu new file mode 100644 index 0000000..4c2d55c --- /dev/null +++ b/fouroversix/src/fouroversix/csrc/quantize/fp4_quant_fp16_nvfp4_trans_sm100.cu @@ -0,0 +1,11 @@ +// Splitting the different transpose modes to different files to speed up compilation. +// This file is auto-generated. See "generate_kernels.py" +#include "fp4_quant_launch_template.h" +namespace fouroversix { + +template<> +void run_fp4_quant_(FP4_quant_params ¶ms, cudaStream_t stream) { + run_nvfp4_quant(params, stream); +} + +} // namespace fouroversix \ No newline at end of file diff --git a/fouroversix/src/fouroversix/matmul/__init__.py b/fouroversix/src/fouroversix/matmul/__init__.py new file mode 100644 index 0000000..7d5807c --- /dev/null +++ b/fouroversix/src/fouroversix/matmul/__init__.py @@ -0,0 +1,3 @@ +from .frontend import fp4_matmul + +__all__ = ["fp4_matmul"] diff --git a/fouroversix/src/fouroversix/matmul/backend.py b/fouroversix/src/fouroversix/matmul/backend.py new file mode 100644 index 0000000..7c932d3 --- /dev/null +++ b/fouroversix/src/fouroversix/matmul/backend.py @@ -0,0 +1,60 @@ +from abc import ABC, abstractmethod + +import torch +from fouroversix.quantize import QuantizedTensor +from fouroversix.utils import DataType + + +class MatmulBackendBase(ABC): + """Base class for all matrix multiplication backends.""" + + @classmethod + @abstractmethod + def is_available(cls) -> bool: + """Return True if the backend is available on the current machine.""" + msg = "Subclasses must implement this method" + raise NotImplementedError(msg) + + @classmethod + @abstractmethod + def is_supported( + cls, + input: QuantizedTensor, + other: QuantizedTensor, + *, + out_dtype: DataType, + ) -> bool: + """Return True if the backend supports the given inputs and output data type.""" + + if not cls.is_available(): + return False + + if input.dtype != other.dtype: + msg = "Both inputs must have the same dtype" + raise ValueError(msg) + + if input.original_shape[1] != other.original_shape[1]: + msg = ( + "The first input must be in row-major layout, the second input must be" + "in column-major layout, and both inputs must have the same inner " + "dimension" + ) + raise ValueError(msg) + + return True + + @classmethod + @abstractmethod + def fp4_matmul( + cls, + input: QuantizedTensor, + other: QuantizedTensor, + *, + out_dtype: DataType, + ) -> torch.Tensor: + """ + Perform a matrix multiplication (`a @ b.T`) between two quantized tensors using + the backend. + """ + msg = "Subclasses must implement this method" + raise NotImplementedError(msg) diff --git a/fouroversix/src/fouroversix/matmul/cutlass/__init__.py b/fouroversix/src/fouroversix/matmul/cutlass/__init__.py new file mode 100644 index 0000000..91f9a43 --- /dev/null +++ b/fouroversix/src/fouroversix/matmul/cutlass/__init__.py @@ -0,0 +1,3 @@ +from .backend import CUTLASSMatmulBackend + +__all__ = ["CUTLASSMatmulBackend"] diff --git a/fouroversix/src/fouroversix/matmul/cutlass/backend.py b/fouroversix/src/fouroversix/matmul/cutlass/backend.py new file mode 100644 index 0000000..7c3d3cc --- /dev/null +++ b/fouroversix/src/fouroversix/matmul/cutlass/backend.py @@ -0,0 +1,147 @@ +import functools + +import torch +from fouroversix.matmul.backend import MatmulBackendBase +from fouroversix.quantize import QuantizedTensor +from fouroversix.utils import BLACKWELL_SM_IDS, SM_100, SM_120, DataType + + +class CUTLASSMatmulBackend(MatmulBackendBase): + """ + The CUTLASS matrix multiplication backend. Uses CUTLASS kernels to perform fast + FP4 matrix multiplication. Requires a Blackwell GPU. + """ + + @classmethod + @functools.lru_cache + def is_available(cls) -> bool: + """Return True if the CUTLASS backend is available on the current machine.""" + + if ( + not torch.cuda.is_available() + or torch.cuda.get_device_capability()[0] not in BLACKWELL_SM_IDS + ): + return False + + try: + import fouroversix._C # noqa: F401 + except ModuleNotFoundError: + return False + + return True + + @classmethod + def is_supported( + cls, + input: QuantizedTensor, + other: QuantizedTensor, + *, + out_dtype: DataType, + ) -> bool: + """ + Return True if the CUTLASS backend supports the given inputs and output data + type. + """ + + if not super().is_supported(input, other, out_dtype=out_dtype): + return False + + return input.device.type == "cuda" + + @classmethod + def fp4_matmul( + cls, + input: QuantizedTensor, + other: QuantizedTensor, + *, + out_dtype: DataType, + ) -> torch.Tensor: + """ + Perform a matrix multiplication (`a @ b.T`) between two quantized tensors using + the CUTLASS backend. + """ + + from .ops import ( + gemm_mxfp4mxfp4_accum_fp32_out_bf16_tnt, + gemm_mxfp4mxfp4_accum_fp32_out_bf16_tnt_sm120, + gemm_nvfp4nvfp4_accum_fp32_out_bf16_tnt, + gemm_nvfp4nvfp4_accum_fp32_out_bf16_tnt_sm120, + gemm_nvfp4nvfp4_accum_fp32_out_fp16_tnt, + gemm_nvfp4nvfp4_accum_fp32_out_fp16_tnt_sm120, + ) + + out_shape = (input.original_shape[0], other.original_shape[0]) + + if input.dtype == DataType.mxfp4: + alpha = torch.ones( + 1, + device=input.values.device, + dtype=torch.float32, + ) + elif input.dtype == DataType.nvfp4: + alpha = ( + (input.amax * other.amax) + / ( + input.scale_rule.max_allowed_e2m1_value() + * input.scale_rule.max_allowed_e4m3_value() + * other.scale_rule.max_allowed_e2m1_value() + * other.scale_rule.max_allowed_e4m3_value() + ) + ).to(torch.float32) + + gemm_fns = { + ( + SM_100, + DataType.mxfp4, + DataType.bfloat16, + ): gemm_mxfp4mxfp4_accum_fp32_out_bf16_tnt, + ( + SM_120, + DataType.mxfp4, + DataType.bfloat16, + ): gemm_mxfp4mxfp4_accum_fp32_out_bf16_tnt_sm120, + ( + SM_100, + DataType.nvfp4, + DataType.bfloat16, + ): gemm_nvfp4nvfp4_accum_fp32_out_bf16_tnt, + ( + SM_120, + DataType.nvfp4, + DataType.bfloat16, + ): gemm_nvfp4nvfp4_accum_fp32_out_bf16_tnt_sm120, + ( + SM_100, + DataType.nvfp4, + DataType.float16, + ): gemm_nvfp4nvfp4_accum_fp32_out_fp16_tnt, + ( + SM_120, + DataType.nvfp4, + DataType.float16, + ): gemm_nvfp4nvfp4_accum_fp32_out_fp16_tnt_sm120, + } + + gemm_fn = gemm_fns.get( + (torch.cuda.get_device_capability()[0], input.dtype, out_dtype), + ) + + if gemm_fn is None: + msg = ( + "No gemm function found for the given device capability and " + f"out_dtype: {torch.cuda.get_device_capability()[0]}, {out_dtype}" + ) + raise ValueError(msg) + + out = gemm_fn( + input.values, + other.values, + input.scale_factors, + other.scale_factors, + alpha, + ) + + if out_shape is not None and out.shape != out_shape: + out = out[: out_shape[0], : out_shape[1]] + + return out diff --git a/fouroversix/src/fouroversix/matmul/cutlass/ops.py b/fouroversix/src/fouroversix/matmul/cutlass/ops.py new file mode 100644 index 0000000..c2ea6ba --- /dev/null +++ b/fouroversix/src/fouroversix/matmul/cutlass/ops.py @@ -0,0 +1,183 @@ +import torch + +import fouroversix._C # noqa: F401 + + +def gemm_mxfp4mxfp4_accum_fp32_out_bf16_tnt( + a: torch.Tensor, + b: torch.Tensor, + a_sf: torch.Tensor, + b_sf: torch.Tensor, + alpha: torch.Tensor, +) -> torch.Tensor: + return torch.ops.fouroversix.gemm_mxfp4mxfp4_accum_fp32_out_bf16_tnt.default( + a, + b, + a_sf, + b_sf, + alpha, + ) + + +@torch.library.register_fake("fouroversix::gemm_mxfp4mxfp4_accum_fp32_out_bf16_tnt") +def _( + a: torch.Tensor, + b: torch.Tensor, + a_sf: torch.Tensor, # noqa: ARG001 + b_sf: torch.Tensor, # noqa: ARG001 + alpha: torch.Tensor, # noqa: ARG001 +) -> torch.Tensor: + m = a.shape[0] + n = b.shape[0] + return torch.empty(m, n, dtype=torch.bfloat16, device=a.device) + + +def gemm_mxfp4mxfp4_accum_fp32_out_bf16_tnt_sm120( + a: torch.Tensor, + b: torch.Tensor, + a_sf: torch.Tensor, + b_sf: torch.Tensor, + alpha: torch.Tensor, +) -> torch.Tensor: + return torch.ops.fouroversix.gemm_mxfp4mxfp4_accum_fp32_out_bf16_tnt_sm120.default( + a, + b, + a_sf, + b_sf, + alpha, + ) + + +@torch.library.register_fake( + "fouroversix::gemm_mxfp4mxfp4_accum_fp32_out_bf16_tnt_sm120", +) +def _( + a: torch.Tensor, + b: torch.Tensor, + a_sf: torch.Tensor, # noqa: ARG001 + b_sf: torch.Tensor, # noqa: ARG001 + alpha: torch.Tensor, # noqa: ARG001 +) -> torch.Tensor: + m = a.shape[0] + n = b.shape[0] + return torch.empty(m, n, dtype=torch.bfloat16, device=a.device) + + +def gemm_nvfp4nvfp4_accum_fp32_out_bf16_tnt( + a: torch.Tensor, + b: torch.Tensor, + a_sf: torch.Tensor, + b_sf: torch.Tensor, + alpha: torch.Tensor, +) -> torch.Tensor: + return torch.ops.fouroversix.gemm_nvfp4nvfp4_accum_fp32_out_bf16_tnt.default( + a, + b, + a_sf, + b_sf, + alpha, + ) + + +@torch.library.register_fake("fouroversix::gemm_nvfp4nvfp4_accum_fp32_out_bf16_tnt") +def _( + a: torch.Tensor, + b: torch.Tensor, + a_sf: torch.Tensor, # noqa: ARG001 + b_sf: torch.Tensor, # noqa: ARG001 + alpha: torch.Tensor, # noqa: ARG001 +) -> torch.Tensor: + m = a.shape[0] + n = b.shape[0] + return torch.empty(m, n, dtype=torch.bfloat16, device=a.device) + + +def gemm_nvfp4nvfp4_accum_fp32_out_bf16_tnt_sm120( + a: torch.Tensor, + b: torch.Tensor, + a_sf: torch.Tensor, + b_sf: torch.Tensor, + alpha: torch.Tensor, +) -> torch.Tensor: + return torch.ops.fouroversix.gemm_nvfp4nvfp4_accum_fp32_out_bf16_tnt_sm120.default( + a, + b, + a_sf, + b_sf, + alpha, + ) + + +@torch.library.register_fake( + "fouroversix::gemm_nvfp4nvfp4_accum_fp32_out_bf16_tnt_sm120", +) +def _( + a: torch.Tensor, + b: torch.Tensor, + a_sf: torch.Tensor, # noqa: ARG001 + b_sf: torch.Tensor, # noqa: ARG001 + alpha: torch.Tensor, # noqa: ARG001 +) -> torch.Tensor: + m = a.shape[0] + n = b.shape[0] + return torch.empty(m, n, dtype=torch.bfloat16, device=a.device) + + +def gemm_nvfp4nvfp4_accum_fp32_out_fp16_tnt( + a: torch.Tensor, + b: torch.Tensor, + a_sf: torch.Tensor, + b_sf: torch.Tensor, + alpha: torch.Tensor, +) -> torch.Tensor: + return torch.ops.fouroversix.gemm_nvfp4nvfp4_accum_fp32_out_fp16_tnt.default( + a, + b, + a_sf, + b_sf, + alpha, + ) + + +@torch.library.register_fake("fouroversix::gemm_nvfp4nvfp4_accum_fp32_out_fp16_tnt") +def _( + a: torch.Tensor, + b: torch.Tensor, + a_sf: torch.Tensor, # noqa: ARG001 + b_sf: torch.Tensor, # noqa: ARG001 + alpha: torch.Tensor, # noqa: ARG001 +) -> torch.Tensor: + m = a.shape[0] + n = b.shape[0] + return torch.empty(m, n, dtype=torch.float16, device=a.device) + + +def gemm_nvfp4nvfp4_accum_fp32_out_fp16_tnt_sm120( + a: torch.Tensor, + b: torch.Tensor, + a_sf: torch.Tensor, + b_sf: torch.Tensor, + alpha: torch.Tensor, +) -> torch.Tensor: + return torch.ops.fouroversix.gemm_nvfp4nvfp4_accum_fp32_out_fp16_tnt_sm120.default( + a, + b, + a_sf, + b_sf, + alpha, + ) + + +@torch.library.register_fake( + "fouroversix::gemm_nvfp4nvfp4_accum_fp32_out_fp16_tnt_sm120", +) +def _( + a: torch.Tensor, + b: torch.Tensor, + a_sf: torch.Tensor, # noqa: ARG001 + b_sf: torch.Tensor, # noqa: ARG001 + alpha: torch.Tensor, # noqa: ARG001 +) -> torch.Tensor: + m = a.shape[0] + n = b.shape[0] + return torch.empty(m, n, dtype=torch.float16, device=a.device) diff --git a/fouroversix/src/fouroversix/matmul/frontend.py b/fouroversix/src/fouroversix/matmul/frontend.py new file mode 100644 index 0000000..bf96062 --- /dev/null +++ b/fouroversix/src/fouroversix/matmul/frontend.py @@ -0,0 +1,115 @@ +import torch +from fouroversix.quantize import QuantizationConfig, QuantizedTensor, quantize_to_fp4 +from fouroversix.utils import DataType, MatmulBackend + +from .cutlass import CUTLASSMatmulBackend +from .pytorch import PyTorchMatmulBackend + +AVAILABLE_BACKENDS = { + MatmulBackend.cutlass: CUTLASSMatmulBackend, + MatmulBackend.pytorch: PyTorchMatmulBackend, +} + + +def fp4_matmul( + input: torch.Tensor | QuantizedTensor, + other: torch.Tensor | QuantizedTensor, + *, + backend: MatmulBackend | None = None, + input_config: QuantizationConfig | None = None, + other_config: QuantizationConfig | None = None, + out_dtype: DataType = DataType.bfloat16, +) -> torch.Tensor: + """ + Perform a matrix multiplication (`a @ b.T`) between two quantized tensors. + + ## Sample Code + + Each tensor may be provided in either high or low precision. If provided in high + precision, tensors will be quantized to FP4 prior to the matrix multiplication, and + quantization may be configured with the `input_quantize_kwargs` and + `other_quantize_kwargs` parameters. For example, the following two code samples are + equivalent: + + ### With High-Precision Inputs + + ```python + a = torch.tensor(1024, 1024, dtype=torch.bfloat16, device="cuda") + b = torch.tensor(1024, 1024, dtype=torch.bfloat16, device="cuda") + out = fp4_matmul(a, b) + ``` + + ### With Low-Precision Inputs + + ```python + a = torch.tensor(1024, 1024, dtype=torch.bfloat16, device="cuda") + b = torch.tensor(1024, 1024, dtype=torch.bfloat16, device="cuda") + + a_quantized = quantize_to_fp4(a) + b_quantized = quantize_to_fp4(b) + + out = fp4_matmul(a_quantized, b_quantized) + ``` + + ## Backends + + We provide two different implementations of FP4 matrix multiplication: + + - **CUTLASS**: Uses CUTLASS kernels to perform fast FP4 matrix multiplication. + Requires a Blackwell GPU. + - **PyTorch**: A slow implementation which dequantizes FP4 tensors and then + performs a high-precision matrix multiplication. + + ## Parameters + + Args: + input (torch.Tensor | QuantizedTensor): The first tensor to be multiplied. + other (torch.Tensor | QuantizedTensor): The second tensor to be multiplied. + backend (MatmulBackend): The backend to use for the matrix multiplication, + either `MatmulBackend.cutlass` or `MatmulBackend.pytorch`. If no backend is + provided, CUTLASS will be used if the machine has a Blackwell GPU, and + PyTorch will be used otherwise. + input_config (QuantizationConfig | None): If `input` is provided in high + precision, this configuration will be passed to the `quantize_to_fp4` call + done prior to the matrix multiplication. + other_config (QuantizationConfig | None): If `other` is provided in high + precision, this configuration will be passed to the `quantize_to_fp4` call + done prior to the matrix multiplication. + out_dtype (DataType): The data type of the output tensor. Defaults to + `DataType.bfloat16`. + + Returns: + The output tensor. + + """ + + if input_config is None: + input_config = QuantizationConfig() + + if isinstance(input, torch.Tensor): + input = quantize_to_fp4(input, input_config) + + if other_config is None: + other_config = QuantizationConfig() + + if isinstance(other, torch.Tensor): + other = quantize_to_fp4(other, other_config) + + if backend is None: + for backend_candidate in [MatmulBackend.cutlass, MatmulBackend.pytorch]: + if AVAILABLE_BACKENDS[backend_candidate] is not None and AVAILABLE_BACKENDS[ + backend_candidate + ].is_supported(input, other, out_dtype=out_dtype): + backend = backend_candidate + break + else: + msg = "No backend found that supports the given parameters" + raise ValueError(msg) + + elif not AVAILABLE_BACKENDS[backend].is_supported( + input, other, out_dtype=out_dtype, + ): + msg = f"Backend {backend} does not support the given parameters" + raise ValueError(msg) + + return AVAILABLE_BACKENDS[backend].fp4_matmul(input, other, out_dtype=out_dtype) diff --git a/fouroversix/src/fouroversix/matmul/pytorch.py b/fouroversix/src/fouroversix/matmul/pytorch.py new file mode 100644 index 0000000..64279c3 --- /dev/null +++ b/fouroversix/src/fouroversix/matmul/pytorch.py @@ -0,0 +1,41 @@ + +import torch +from fouroversix.quantize import QuantizedTensor +from fouroversix.utils import DataType + +from .backend import MatmulBackendBase + + +class PyTorchMatmulBackend(MatmulBackendBase): + """ + The PyTorch matrix multiplication backend. Dequantizes both inputs to FP32 and + performs an FP32 matrix multiplication in order to simulate an NVFP4 matrix + multiplication which accumulates in FP32. Slow, but can be run on any GPU. + """ + + @classmethod + def is_available(cls) -> bool: + """Return True if the PyTorch backend is available on the current machine.""" + return True + + @classmethod + def fp4_matmul( + cls, + input: QuantizedTensor, + other: QuantizedTensor, + *, + out_dtype: DataType, + ) -> torch.Tensor: + """Perform a matrix multiplication (`a @ b.T`) between two quantized tensors.""" + + out_shape = (input.original_shape[0], other.original_shape[0]) + + out = torch.matmul( + input.dequantize(dtype=torch.float32), + other.dequantize(dtype=torch.float32).T, + ).to(out_dtype.torch_dtype()) + + if out.shape != out_shape: + out = out[: out_shape[0], : out_shape[1]] + + return out diff --git a/fouroversix/src/fouroversix/model/__init__.py b/fouroversix/src/fouroversix/model/__init__.py new file mode 100644 index 0000000..7062799 --- /dev/null +++ b/fouroversix/src/fouroversix/model/__init__.py @@ -0,0 +1,11 @@ +from .config import ModelQuantizationConfig, ModuleQuantizationConfig +from .modules import FourOverSixLinear +from .quantize import QuantizedModule, quantize_model + +__all__ = [ + "FourOverSixLinear", + "ModelQuantizationConfig", + "ModuleQuantizationConfig", + "QuantizedModule", + "quantize_model", +] diff --git a/fouroversix/src/fouroversix/model/config.py b/fouroversix/src/fouroversix/model/config.py new file mode 100644 index 0000000..3c8b5fb --- /dev/null +++ b/fouroversix/src/fouroversix/model/config.py @@ -0,0 +1,194 @@ +import hashlib +import json +from dataclasses import dataclass, field +from typing import Any + +from fouroversix.quantize import QuantizationConfig +from fouroversix.utils import ( + DataType, + MatmulBackend, + QuantizeBackend, + RoundStyle, + ScaleRule, +) + + +@dataclass +class ModuleQuantizationConfig: + """ + Configuration for quantizing modules with Four Over Six. + + Args: + activation_scale_rule (ScaleRule | None): The scaling rule to use for activation + tensors. If not provided, `scale_rule` will be used. + dtype (DataType): The quantization data type to use for the module. Defaults to + `DataType.nvfp4`. + gradient_round_style (RoundStyle | None): The rounding style to use for gradient + tensors. Defaults to `RoundStyle.stochastic`. + gradient_scale_rule (ScaleRule | None): The scaling rule to use for gradient + tensors. If not provided, `scale_rule` will be used. + keep_master_weights (bool): Whether to keep the master weights. Defaults to + `False`. + matmul_backend (MatmulBackend | None): The backend to use for matrix + multiplications. If not provided, a backend will be selected automatically + based on the available GPU and the specified options. + output_dtype (DataType): The data type to use for the module's output. Defaults + to `DataType.bfloat16`. + quantize_backend (QuantizeBackend | None): The backend to use for quantization. + If not provided, a backend will be selected automatically based on the + available GPU and the specified options. + scale_rule (ScaleRule): The fallback scaling rule which will be used if any of + the other scaling rules are not specified. + weight_scale_2d (bool): Whether to use 2D block scaling for weights. Should be + set to `True` if the module is used for training. + weight_scale_rule (ScaleRule | None): The scaling rule to use for weights. If + not provided, `scale_rule` will be used. + + """ + + activation_scale_rule: ScaleRule | None = None + dtype: DataType = DataType.nvfp4 + gradient_round_style: RoundStyle = RoundStyle.stochastic + gradient_scale_rule: ScaleRule | None = None + keep_master_weights: bool = False + matmul_backend: MatmulBackend | None = None + output_dtype: DataType = DataType.bfloat16 + quantize_backend: QuantizeBackend | None = None + scale_rule: ScaleRule = ScaleRule.mse + weight_scale_2d: bool = False + weight_scale_rule: ScaleRule | None = None + activation_chunk_size: int | None = None + + def __post_init__(self) -> None: + """Convert string values to enums.""" + + if isinstance(self.activation_scale_rule, str): + self.activation_scale_rule = ScaleRule(self.activation_scale_rule) + + if isinstance(self.dtype, str): + self.dtype = DataType(self.dtype) + + if isinstance(self.gradient_round_style, str): + self.gradient_round_style = RoundStyle(self.gradient_round_style) + + if isinstance(self.gradient_scale_rule, str): + self.gradient_scale_rule = ScaleRule(self.gradient_scale_rule) + + if isinstance(self.matmul_backend, str): + self.matmul_backend = MatmulBackend(self.matmul_backend) + + if isinstance(self.output_dtype, str): + self.output_dtype = DataType(self.output_dtype) + + if isinstance(self.quantize_backend, str): + self.quantize_backend = QuantizeBackend(self.quantize_backend) + + if isinstance(self.scale_rule, str): + self.scale_rule = ScaleRule(self.scale_rule) + + if isinstance(self.weight_scale_rule, str): + self.weight_scale_rule = ScaleRule(self.weight_scale_rule) + + self.activation_scale_rule = self.activation_scale_rule or self.scale_rule + self.gradient_scale_rule = self.gradient_scale_rule or self.scale_rule + self.weight_scale_rule = self.weight_scale_rule or self.scale_rule + + def get_activation_config(self, **kwargs: dict[str, Any]) -> QuantizationConfig: + """Return the quantization configuration for the activation tensors.""" + return QuantizationConfig( + backend=self.quantize_backend, + dtype=self.dtype, + scale_rule=self.activation_scale_rule, + **kwargs, + ) + + def get_gradient_config(self, **kwargs: dict[str, Any]) -> QuantizationConfig: + """Return the quantization configuration for the gradient tensors.""" + return QuantizationConfig( + backend=self.quantize_backend, + dtype=self.dtype, + round_style=self.gradient_round_style, + scale_rule=self.gradient_scale_rule, + **kwargs, + ) + + def get_weight_config(self, **kwargs: dict[str, Any]) -> QuantizationConfig: + """Return the quantization configuration for the weight tensors.""" + return QuantizationConfig( + backend=self.quantize_backend, + block_scale_2d=self.weight_scale_2d, + dtype=self.dtype, + scale_rule=self.weight_scale_rule, + **kwargs, + ) + + +@dataclass +class ModelQuantizationConfig(ModuleQuantizationConfig): + """ + Configuration for quantizing a model with Four Over Six. + + Args: + activation_scale_rule (ScaleRule | None): The scaling rule to use for activation + tensors. If not provided, `scale_rule` will be used. + dtype (DataType): The quantization data type to use for the module. Defaults to + `DataType.nvfp4`. + gradient_round_style (RoundStyle | None): The rounding style to use for gradient + tensors. Defaults to `RoundStyle.stochastic`. + gradient_scale_rule (ScaleRule | None): The scaling rule to use for gradient + tensors. If not provided, `scale_rule` will be used. + keep_master_weights (bool): Whether to keep the master weights. Defaults to + `False`. + matmul_backend (MatmulBackend | None): The backend to use for matrix + multiplications. If not provided, a backend will be selected automatically + based on the available GPU and the specified options. + output_dtype (DataType): The data type to use for the module's output. Defaults + to `DataType.bfloat16`. + quantize_backend (QuantizeBackend | None): The backend to use for quantization. + If not provided, a backend will be selected automatically based on the + available GPU and the specified options. + scale_rule (ScaleRule): The fallback scaling rule which will be used if any of + the other scaling rules are not specified. + weight_scale_2d (bool): Whether to use 2D block scaling for weights. Should be + set to `True` if the module is used for training. + weight_scale_rule (ScaleRule | None): The scaling rule to use for weights. If + not provided, `scale_rule` will be used. + + module_config_overrides (dict[str, ModuleQuantizationConfig]): A mapping of + module names to quantization configurations to use for each module. If a + module is not specified, the attributes from this class will be used. + modules_to_not_convert (list[str]): A list of module names that should not be + quantized. + + """ + + module_config_overrides: dict[str, ModuleQuantizationConfig] = field( + default_factory=dict, + ) + modules_to_not_convert: list[str] = field(default_factory=lambda: ["lm_head"]) + + def __post_init__(self) -> None: + """Convert module config overrides to ModuleQuantizationConfig instances.""" + + super().__post_init__() + + if self.module_config_overrides is not None: + for module_name, module_config in self.module_config_overrides.items(): + if isinstance(module_config, dict): + self.module_config_overrides[module_name] = ( + ModuleQuantizationConfig(**module_config) + ) + + def get_module_config(self, module_name: str) -> ModuleQuantizationConfig: + """Return the quantization configuration for a given module.""" + return ( + self.module_config_overrides.get(module_name, self) + if self.module_config_overrides is not None + else self + ) + + def __hash__(self) -> str: + """Return a hash of the configuration.""" + return hashlib.sha256( + json.dumps(self.__dict__, sort_keys=True).encode(), + ).hexdigest() diff --git a/fouroversix/src/fouroversix/model/modules/__init__.py b/fouroversix/src/fouroversix/model/modules/__init__.py new file mode 100644 index 0000000..6bda99f --- /dev/null +++ b/fouroversix/src/fouroversix/model/modules/__init__.py @@ -0,0 +1,7 @@ +from .gpt_oss import FourOverSixGptOssMLP +from .linear import FourOverSixLinear +# from .qwen import FourOverSixQwenExperts + +# __all__ = ["FourOverSixGptOssMLP", "FourOverSixLinear", "FourOverSixQwenExperts"] +__all__ = ["FourOverSixGptOssMLP", "FourOverSixLinear"] + diff --git a/fouroversix/src/fouroversix/model/modules/gpt_oss.py b/fouroversix/src/fouroversix/model/modules/gpt_oss.py new file mode 100644 index 0000000..b181fa5 --- /dev/null +++ b/fouroversix/src/fouroversix/model/modules/gpt_oss.py @@ -0,0 +1,395 @@ +from typing import Any + +import torch +from fouroversix.matmul import fp4_matmul +from fouroversix.model.config import ModuleQuantizationConfig +from fouroversix.model.quantize import QuantizedModule +from fouroversix.quantize import ( + QuantizationConfig, + QuantizedTensor, + quantize_to_fp4, +) +from torch import nn +from transformers import GptOssConfig +from transformers.models.gpt_oss.modeling_gpt_oss import ( + GptOssExperts, + GptOssMLP, + GptOssTopKRouter, +) + + +@QuantizedModule.register(GptOssMLP) +class FourOverSixGptOssMLP(nn.Module): + """Drop-in replacement for GptOssMLP layer that uses FP4 quantization.""" + + def __init__( + self, + module: GptOssMLP, + config: ModuleQuantizationConfig, + ) -> None: + """ + Initialize the FourOverSixGptOssMLP layer. + + Args: + module (GptOssMLP): The high-precision module that this quantized layer will + replace. + config (ModuleQuantizationConfig): The quantization configuration to use for + the layer. + + """ + + super().__init__() + + self.config = config + + gpt_oss_config = GptOssConfig( + num_local_experts=module.experts.num_experts, + hidden_size=module.experts.hidden_size, + intermediate_size=module.experts.intermediate_size, + num_experts_per_token=module.router.top_k, + ) + + self.router = GptOssTopKRouter(gpt_oss_config) + self.router.weight = module.router.weight + self.router.bias = module.router.bias + + self.experts = FourOverSixGptOssExperts( + module.experts, + quantization_config=self.config, + ) + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Forward pass for the FP4 MLP layer.""" + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_states = hidden_states.reshape(-1, hidden_dim) + _, router_scores, router_indices = self.router(hidden_states) + hidden_states = self.experts(hidden_states, router_indices, router_scores) + hidden_states = hidden_states.reshape(batch_size, sequence_length, hidden_dim) + return hidden_states, router_scores + + @property + def parameters_to_quantize(self) -> tuple[str, ...]: + """Return high precision parameters to be quantized and deleted.""" + return () + + def get_element_size(self, parameter_name: str) -> float: + """Get the size of a single element, in bytes, for a parameter.""" + + return { + "quantized_down_proj_values": 0.5, + "quantized_gate_up_proj_values": 0.5, + "down_proj": 9 / 16, + "gate_up_proj": 9 / 16, + }.get( + parameter_name, + getattr(self, parameter_name).element_size(), + ) + + +@QuantizedModule.register(GptOssExperts, replace_existing_modules_in_model=False) +class FourOverSixGptOssExperts(nn.Module): + """Drop-in replacement for GptOssExperts layer that uses FP4 quantization.""" + + def __init__( + self, + module: GptOssExperts, + quantization_config: ModuleQuantizationConfig | None = None, + ) -> None: + + super().__init__() + + self.num_experts = module.num_experts + self.intermediate_size = module.intermediate_size + self.hidden_size = module.hidden_size + + self.down_proj_bias = module.down_proj_bias + self.gate_up_proj_bias = module.gate_up_proj_bias + + # Store original weights so they can be quantized and then deleted + self.down_proj = module.down_proj + self.gate_up_proj = module.gate_up_proj + + self.config = quantization_config + + if not self.config.keep_master_weights: + self.register_buffer( + "quantized_down_proj_values", + nn.Parameter( + torch.zeros( + self.num_experts, + self.intermediate_size, + self.hidden_size // 2, + dtype=torch.uint8, + ), + requires_grad=False, + ), + ) + + self.register_buffer( + "quantized_gate_up_proj_values", + nn.Parameter( + torch.zeros( + self.num_experts, + self.intermediate_size * 2, + self.hidden_size // 2, + dtype=torch.uint8, + ), + requires_grad=False, + ), + ) + + self.register_buffer( + "quantized_down_proj_scale_factors", + nn.Parameter( + torch.zeros( + self.num_experts, + self.hidden_size + * self.intermediate_size + // self.config.dtype.block_size(), + dtype=self.config.dtype.scale_dtype(), + ), + requires_grad=False, + ), + ) + self.register_buffer( + "quantized_gate_up_proj_scale_factors", + nn.Parameter( + torch.zeros( + self.num_experts, + self.hidden_size + * (self.intermediate_size * 2) + // self.config.dtype.block_size(), + dtype=self.config.dtype.scale_dtype(), + ), + requires_grad=False, + ), + ) + + self.register_buffer( + "quantized_down_proj_amax", + nn.Parameter( + torch.zeros(self.num_experts, 1, dtype=torch.float32), + requires_grad=False, + ), + ) + self.register_buffer( + "quantized_gate_up_proj_amax", + nn.Parameter( + torch.zeros(self.num_experts, 1, dtype=torch.float32), + requires_grad=False, + ), + ) + + self.register_buffer( + "quantized_down_proj_metadata", + nn.Parameter( + torch.zeros(self.num_experts, 4, dtype=torch.int32), + requires_grad=False, + ), + ) + self.register_buffer( + "quantized_gate_up_proj_metadata", + nn.Parameter( + torch.zeros(self.num_experts, 4, dtype=torch.int32), + requires_grad=False, + ), + ) + + self.alpha = 1.702 + self.limit = 7.0 + + @property + def parameters_to_quantize(self) -> tuple[str, ...]: + """Return high precision parameters to be quantized and deleted.""" + return ("down_proj", "gate_up_proj") + + def get_packing_factor(self, parameter_name: str) -> float: + """Get the packing factor for a parameter.""" + return ( + 2 + if parameter_name + in {"quantized_down_proj_values", "quantized_gate_up_proj_values"} + else 1 + ) + + def get_quantized_parameters( + self, + parameter_name: str, + parameter: torch.Tensor, + ) -> dict[str, Any]: + """ + Prepare this layer for post-training quantization by quantizing the weight, + storing the quantized weight, and deleting the original weight. This should not + be done if the layer is used for training, as training requires storage of the + high-precision weight. + """ + + weight_config = QuantizationConfig( + backend=self.config.quantize_backend, + dtype=self.config.dtype, + scale_rule=self.config.weight_scale_rule, + ) + + quantized_proj = [] + for e in range(parameter.shape[0]): + q = quantize_to_fp4(parameter[e], weight_config) + quantized_proj.append(q) + + if "down" in parameter_name: + prefix = "down" + elif "gate_up" in parameter_name: + prefix = "gate_up" + + return { + f"quantized_{prefix}_proj_values": torch.stack( + [tensor.values for tensor in quantized_proj], + dim=0, + ), + f"quantized_{prefix}_proj_scale_factors": torch.stack( + [tensor.scale_factors for tensor in quantized_proj], + dim=0, + ), + f"quantized_{prefix}_proj_amax": torch.stack( + [tensor.amax for tensor in quantized_proj], + dim=0, + ), + f"quantized_{prefix}_proj_metadata": torch.stack( + [ + torch.tensor( + [ + tensor.original_shape[0], + tensor.original_shape[1], + tensor.padded_shape[0], + tensor.padded_shape[1], + ], + ) + for tensor in quantized_proj + ], + ), + } + + def forward( + self, + hidden_states: torch.Tensor, + routing_indices: torch.Tensor = None, + routing_weights: torch.Tensor = None, + ) -> torch.Tensor: + """Forward pass for the FP4 experts layer.""" + + down_proj, gate_up_proj = self.quantized_weights() + + batch_size = hidden_states.shape[0] + hidden_states = hidden_states.reshape( + -1, + self.hidden_size, + ) + next_states = torch.zeros_like( + hidden_states, + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + with torch.no_grad(): + expert_mask = torch.nn.functional.one_hot( + routing_indices, + num_classes=self.num_experts, + ) + expert_mask = expert_mask.permute(2, 1, 0) + expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() + + for [expert_idx] in expert_hit: + if expert_idx == self.num_experts: + continue + with torch.no_grad(): + top_k_pos, token_idx = torch.where(expert_mask[expert_idx]) + current_state = hidden_states[token_idx] + + # Gate-up projection + fprop_activation_config = QuantizationConfig( + backend=self.config.quantize_backend, + dtype=self.config.dtype, + scale_rule=self.config.activation_scale_rule, + ) + + gate_up = fp4_matmul( + current_state, + gate_up_proj[expert_idx], + input_config=fprop_activation_config, + out_dtype=self.config.output_dtype, + ) + gate_up += self.gate_up_proj_bias[expert_idx] + + gate, up = gate_up[..., ::2], gate_up[..., 1::2] + gate = gate.clamp(min=None, max=self.limit) + up = up.clamp(min=-self.limit, max=self.limit) + glu = gate * torch.sigmoid(gate * self.alpha) + gated_output = (up + 1) * glu + + # Down projection + out = fp4_matmul( + gated_output, + down_proj[expert_idx], + input_config=fprop_activation_config, + out_dtype=self.config.output_dtype, + ) + out += self.down_proj_bias[expert_idx] + weighted_output = out * routing_weights[token_idx, top_k_pos, None] + next_states.index_add_( + 0, + token_idx, + weighted_output.to(hidden_states.dtype), + ) + + return next_states.view(batch_size, -1, self.hidden_size) + + def quantized_weights(self) -> tuple[list[QuantizedTensor], list[QuantizedTensor]]: + """Return quantized parameters as QuantizedTensor.""" + + if not hasattr(self, "_quantized_weights"): + weight_config = self.config.get_weight_config() + if self.config.keep_master_weights: + down = [ + quantize_to_fp4(self.down_proj[e], weight_config) + for e in range(self.num_experts) + ] + gate_up = [ + quantize_to_fp4(self.gate_up_proj[e], weight_config) + for e in range(self.num_experts) + ] + return (down, gate_up) + + down = [] + gate_up = [] + for e in range(self.num_experts): + down.append( + QuantizedTensor( + values=self.quantized_down_proj_values.data[e], + scale_factors=self.quantized_down_proj_scale_factors.data[e], + amax=self.quantized_down_proj_amax.data[e], + dtype=self.config.dtype, + original_shape=tuple( + self.quantized_down_proj_metadata.data[e, :2].tolist(), + ), + scale_rule=self.config.weight_scale_rule, + padded_shape=tuple( + self.quantized_down_proj_metadata.data[e, 2:].tolist(), + ), + ), + ) + gate_up.append( + QuantizedTensor( + values=self.quantized_gate_up_proj_values.data[e], + scale_factors=self.quantized_gate_up_proj_scale_factors.data[e], + amax=self.quantized_gate_up_proj_amax.data[e], + dtype=self.config.dtype, + original_shape=tuple( + self.quantized_gate_up_proj_metadata.data[e, :2].tolist(), + ), + scale_rule=self.config.weight_scale_rule, + padded_shape=tuple( + self.quantized_gate_up_proj_metadata.data[e, 2:].tolist(), + ), + ), + ) + self._quantized_weights = (down, gate_up) + + return self._quantized_weights diff --git a/fouroversix/src/fouroversix/model/modules/linear.py b/fouroversix/src/fouroversix/model/modules/linear.py new file mode 100644 index 0000000..9cf2e1e --- /dev/null +++ b/fouroversix/src/fouroversix/model/modules/linear.py @@ -0,0 +1,355 @@ +from typing import Any +import torch +import torch.nn as nn +from fouroversix.matmul import fp4_matmul +from fouroversix.model.config import ModuleQuantizationConfig +from fouroversix.model.quantize import QuantizedModule +from fouroversix.quantize import ( + QuantizationConfig, + QuantizedTensor, + quantize_to_fp4, + ) + + +class FourOverSixLinearFunction(torch.autograd.Function): + + """Differentiable FP4 linear layer.""" + @staticmethod + def forward( + ctx: torch.autograd.function.FunctionCtx, + config: ModuleQuantizationConfig, + input: torch.Tensor, + weight: torch.Tensor | QuantizedTensor, + weight_t: QuantizedTensor | None, + bias: torch.Tensor = None, + ) -> tuple[torch.Tensor,]: + """ + Perform an FP4 matrix multiplication. The input is provided in high precision + and quantized to FP4 prior to the matrix multiplication, while the weight is + provided in low precision. + """ + needs_wgrad = isinstance(weight, (nn.Parameter, torch.Tensor)) and weight.requires_grad + needs_input_grad = input.requires_grad + ctx.config = config + ctx.needs_wgrad = needs_wgrad + ctx.needs_input_grad_flag = needs_input_grad + ctx.saved_quantized_weight = None + ctx.saved_quantized_weight_t = None + ctx.saved_bias = None + fprop_activation_config = config.get_activation_config() + fprop_weight_config = config.get_weight_config() + if isinstance(weight, QuantizedTensor): + weight_q = weight + else: + weight_q = quantize_to_fp4(weight.data if isinstance(weight, nn.Parameter) else weight, fprop_weight_config) + input_2d = input.reshape(-1, input.shape[-1]) + input_q = quantize_to_fp4(input_2d, fprop_activation_config) + if needs_wgrad: + ctx.save_for_backward( + input_q.values, input_q.scale_factors, input_q.amax, + weight, bias, + ) + ctx.input_shape = input.shape + ctx.input_q_meta = ( + input_q.original_shape, input_q.padded_shape, + input_q.dtype, input_q.scale_rule, + ) + elif needs_input_grad: + # When master weights are dropped, `weight` is a QuantizedTensor rather than + # a Tensor/Parameter, so it cannot be passed to save_for_backward. + if isinstance(weight, QuantizedTensor): + ctx.saved_quantized_weight = weight + ctx.saved_quantized_weight_t = weight_t + ctx.saved_bias = bias + else: + ctx.save_for_backward(weight, bias) + ctx.input_shape = input.shape + out = fp4_matmul( + input_q, + weight_q, + backend=config.matmul_backend, + out_dtype=config.output_dtype, + ).reshape(*input.shape[:-1], weight_q.original_shape[0]) + if bias is not None: + out = out + bias + return out + + @staticmethod + def backward( + ctx: torch.autograd.function.FunctionCtx, + grad_output: torch.Tensor, + ) -> tuple[torch.Tensor, ...]: + """Backward pass for the FP4 linear layer.""" + if not ctx.needs_wgrad and not ctx.needs_input_grad_flag: + return None, None, None, None, None + + if ctx.needs_wgrad: + iq_vals, iq_sf, iq_amax, weight, bias = ctx.saved_tensors + weight_t = None + else: + if ctx.saved_quantized_weight is not None: + weight = ctx.saved_quantized_weight + weight_t = ctx.saved_quantized_weight_t + bias = ctx.saved_bias + else: + weight, bias = ctx.saved_tensors + weight_t = None + + input_shape = ctx.input_shape + + grad_input = None + if ctx.needs_input_grad_flag: + dgrad_grad_config = ctx.config.get_gradient_config() + dgrad_weight_config = ctx.config.get_weight_config(transpose=True) + dgrad_weight = weight_t if weight_t is not None else weight + if isinstance(weight, QuantizedTensor) and weight_t is None: + raise RuntimeError( + "Materialized quantized linear layers need cached transposed weights for grad_input. " + "Re-materialize with cache_transposed_weights=True." + ) + grad_input = fp4_matmul( + grad_output.reshape(-1, grad_output.shape[-1]), + dgrad_weight, + backend=ctx.config.matmul_backend, + input_config=dgrad_grad_config, + other_config=dgrad_weight_config, + out_dtype=ctx.config.output_dtype, + ).reshape(input_shape) + + grad_weight = None + if ctx.needs_wgrad: + orig_shape, padded_shape, q_dtype, q_scale_rule = ctx.input_q_meta + input_approx = QuantizedTensor( + iq_vals, iq_sf, iq_amax, + q_dtype, orig_shape, q_scale_rule, padded_shape, + ).dequantize_triton(dtype=torch.bfloat16) + + wgrad_grad_config = ctx.config.get_gradient_config(rht=True, transpose=True) + wgrad_activation_config = ctx.config.get_activation_config( + rht=True, + transpose=True, + ) + grad_weight = fp4_matmul( + grad_output.reshape(-1, grad_output.shape[-1]), + input_approx, + backend=ctx.config.matmul_backend, + input_config=wgrad_grad_config, + other_config=wgrad_activation_config, + out_dtype=ctx.config.output_dtype, + ).unsqueeze(0) + + grad_bias = ( + grad_output.sum(0) if bias is not None and ctx.needs_input_grad[4] else None + ) + + return ( + None, + grad_input, + grad_weight, + None, + grad_bias, + ) + +@QuantizedModule.register(nn.Linear) +class FourOverSixLinear(nn.Linear): + """ + Drop-in replacement for `nn.Linear` that quantizes weights, activations, and + gradients. + """ + def __init__( + self, + module: nn.Linear, + config: ModuleQuantizationConfig, + ) -> None: + """ + Initialize the FourOverSixLinear layer. + Args: + module (nn.Linear): The high-precision module that this quantized layer will + replace. + config (ModuleQuantizationConfig): The quantization configuration to use for + the layer. + """ + super().__init__( + module.in_features, + module.out_features, + module.bias is not None, + module.weight.device, + module.weight.dtype, + ) + self.weight = module.weight + self.bias = module.bias + self.config = config + if not self.config.keep_master_weights: + self.register_buffer( + "quantized_weight_values", + nn.Parameter( + torch.zeros( + self.out_features, + self.in_features // 2, + dtype=torch.uint8, + ), + requires_grad=False, + ), + ) + self.register_buffer( + "quantized_weight_scale_factors", + nn.Parameter( + torch.zeros( + self.out_features + * self.in_features + // self.config.dtype.block_size(), + dtype=self.config.dtype.scale_dtype(), + ), + requires_grad=False, + ), + ) + self.register_buffer( + "quantized_weight_amax", + nn.Parameter(torch.zeros(1, dtype=torch.float32), requires_grad=False), + ) + self.register_buffer( + "quantized_weight_metadata", + nn.Parameter( + torch.zeros(2 + 2, dtype=torch.int32), + requires_grad=False, + ), + ) + + @property + def parameters_to_quantize(self) -> tuple[str, ...]: + """Return high precision parameters to be quantized and deleted.""" + return ("weight",) + + def get_element_size(self, parameter_name: str) -> float: + """Get the size of a single element, in bytes, for a parameter.""" + # quantized_weight_values is packed, so there are 4 bits, or 0.5 bytes, per + # element. Once quantized, weight will have (8+1)/16 bytes per element (one + # block of 16 values is 8 bytes of values + 1 byte of scale factors). + return {"quantized_weight_values": 0.5, "weight": 9 / 16}.get( + parameter_name, + getattr(self, parameter_name).element_size(), + ) + + def get_quantized_parameters( + self, + parameter_name: str, + parameter: torch.Tensor, + include_transposed: bool = False, + ) -> dict[str, Any]: + """Get the quantized parameters for the layer.""" + if parameter_name == "weight": + config = QuantizationConfig( + backend=self.config.quantize_backend, + block_scale_2d=self.config.weight_scale_2d, + dtype=self.config.dtype, + scale_rule=self.config.weight_scale_rule, + ) + quantized_weight = quantize_to_fp4(parameter, config) + quantized_params = self._serialize_quantized_weight( + "quantized_weight", + quantized_weight, + ) + if include_transposed: + transposed_config = QuantizationConfig( + backend=self.config.quantize_backend, + block_scale_2d=self.config.weight_scale_2d, + dtype=self.config.dtype, + scale_rule=self.config.weight_scale_rule, + transpose=True, + ) + quantized_params.update( + self._serialize_quantized_weight( + "quantized_weight_transposed", + quantize_to_fp4(parameter, transposed_config), + ), + ) + return quantized_params + msg = f"Unsupported high-preciison parameter: {parameter_name}" + raise ValueError(msg) + + @staticmethod + def _serialize_quantized_weight( + prefix: str, + quantized_weight: QuantizedTensor, + ) -> dict[str, torch.Tensor]: + return { + f"{prefix}_values": quantized_weight.values, + f"{prefix}_scale_factors": quantized_weight.scale_factors, + f"{prefix}_amax": quantized_weight.amax, + f"{prefix}_metadata": torch.tensor( + [ + quantized_weight.original_shape[0], + quantized_weight.original_shape[1], + quantized_weight.padded_shape[0], + quantized_weight.padded_shape[1], + ], + dtype=torch.int32, + ), + } + + def _build_quantized_weight(self, prefix: str) -> QuantizedTensor | None: + values = getattr(self, f"{prefix}_values", None) + scale_factors = getattr(self, f"{prefix}_scale_factors", None) + amax = getattr(self, f"{prefix}_amax", None) + metadata = getattr(self, f"{prefix}_metadata", None) + if any(x is None for x in (values, scale_factors, amax, metadata)): + return None + original_shape = tuple(metadata.data[:2].tolist()) + padded_shape = tuple(metadata.data[2:].tolist()) + return QuantizedTensor( + values.data, + scale_factors.data, + amax.data, + self.config.dtype, + original_shape, + self.config.weight_scale_rule, + padded_shape, + ) + + def quantized_weight(self) -> QuantizedTensor: + """ + Prepare this layer for post-training quantization by quantizing the weight, + storing the quantized weight, and deleting the original weight. This should not + be done if the layer is used for training, as training requires storage of the + high-precision weight. + """ + if not hasattr(self, "_quantized_weight"): + if self.config.keep_master_weights: + return self.weight + self._quantized_weight = self._build_quantized_weight("quantized_weight") + return self._quantized_weight + + def quantized_weight_transposed(self) -> QuantizedTensor | None: + if self.config.keep_master_weights: + return None + if not hasattr(self, "_quantized_weight_transposed"): + self._quantized_weight_transposed = self._build_quantized_weight( + "quantized_weight_transposed", + ) + return self._quantized_weight_transposed + + _INT32_MAX = 2_147_483_647 + + def forward(self, input: torch.Tensor) -> torch.Tensor: + """Forward pass for the FP4 linear layer.""" + flat_rows = input.reshape(-1, input.shape[-1]).shape[0] + if flat_rows * input.shape[-1] > self._INT32_MAX: + return self._forward_chunked(input) + return FourOverSixLinearFunction.apply( + self.config, + input, + self.quantized_weight(), + self.quantized_weight_transposed(), + self.bias, + ) + + def _forward_chunked(self, input: torch.Tensor) -> torch.Tensor: + """Split evenly into 2 chunks to avoid int32 overflow in CUDA FP4 quantize kernel.""" + orig_shape = input.shape + input_2d = input.reshape(-1, input.shape[-1]) + mid = input_2d.shape[0] // 2 + weight = self.quantized_weight() + weight_t = self.quantized_weight_transposed() + out_a = FourOverSixLinearFunction.apply(self.config, input_2d[:mid], weight, weight_t, self.bias) + out_b = FourOverSixLinearFunction.apply(self.config, input_2d[mid:], weight, weight_t, self.bias) + return torch.cat([out_a, out_b], dim=0).reshape(*orig_shape[:-1], -1) \ No newline at end of file diff --git a/fouroversix/src/fouroversix/model/modules/linear_ori.py b/fouroversix/src/fouroversix/model/modules/linear_ori.py new file mode 100644 index 0000000..0d88a07 --- /dev/null +++ b/fouroversix/src/fouroversix/model/modules/linear_ori.py @@ -0,0 +1,257 @@ +from typing import Any +import torch +import torch.nn as nn +from fouroversix.matmul import fp4_matmul +from fouroversix.model.config import ModuleQuantizationConfig +from fouroversix.model.quantize import QuantizedModule +from fouroversix.quantize import ( +QuantizationConfig, +QuantizedTensor, +quantize_to_fp4, +) + + +class FourOverSixLinearFunction(torch.autograd.Function): + + """Differentiable FP4 linear layer.""" + @staticmethod + def forward( + ctx: torch.autograd.function.FunctionCtx, + config: ModuleQuantizationConfig, + input: torch.Tensor, + weight: torch.Tensor | QuantizedTensor, + bias: torch.Tensor = None, + ) -> tuple[torch.Tensor,]: + """ + Perform an FP4 matrix multiplication. The input is provided in high precision + and quantized to FP4 prior to the matrix multiplication, while the weight is + provided in low precision. + """ + needs_wgrad = isinstance(weight, (nn.Parameter, torch.Tensor)) and weight.requires_grad + ctx.config = config + ctx.needs_wgrad = needs_wgrad + fprop_activation_config = config.get_activation_config() + fprop_weight_config = config.get_weight_config() + if isinstance(weight, QuantizedTensor): + weight_q = weight + else: + weight_q = quantize_to_fp4(weight.data if isinstance(weight, nn.Parameter) else weight, fprop_weight_config) + input_2d = input.reshape(-1, input.shape[-1]) + input_q = quantize_to_fp4(input_2d, fprop_activation_config) + if needs_wgrad: + # Save FP4 components (~4x smaller than BF16) instead of full-precision input + ctx.save_for_backward( + input_q.values, input_q.scale_factors, input_q.amax, + weight, bias, + ) + ctx.input_shape = input.shape + ctx.input_q_meta = ( + input_q.original_shape, input_q.padded_shape, + input_q.dtype, input_q.scale_rule, + ) + out = fp4_matmul( + input_q, + weight_q, + backend=config.matmul_backend, + out_dtype=config.output_dtype, + ).reshape(*input.shape[:-1], weight_q.original_shape[0]) + if bias is not None: + out = out + bias + return out + + @staticmethod + def backward( + ctx: torch.autograd.function.FunctionCtx, + grad_output: torch.Tensor, + ) -> tuple[torch.Tensor, ...]: + """Backward pass for the FP4 linear layer.""" + if not ctx.needs_wgrad: + return None, None, None, None + iq_vals, iq_sf, iq_amax, weight, bias = ctx.saved_tensors + input_shape = ctx.input_shape + + orig_shape, padded_shape, q_dtype, q_scale_rule = ctx.input_q_meta + input_approx = QuantizedTensor( + iq_vals, iq_sf, iq_amax, + q_dtype, orig_shape, q_scale_rule, padded_shape, + ).dequantize_triton(dtype=torch.bfloat16) + + dgrad_grad_config = ctx.config.get_gradient_config() + dgrad_weight_config = ctx.config.get_weight_config(transpose=True) + + grad_input = fp4_matmul( + grad_output.reshape(-1, grad_output.shape[-1]), + weight, + backend=ctx.config.matmul_backend, + input_config=dgrad_grad_config, + other_config=dgrad_weight_config, + out_dtype=ctx.config.output_dtype, + ).reshape(input_shape) + + wgrad_grad_config = ctx.config.get_gradient_config(rht=True, transpose=True) + wgrad_activation_config = ctx.config.get_activation_config( + rht=True, + transpose=True, + ) + + grad_weight = fp4_matmul( + grad_output.reshape(-1, grad_output.shape[-1]), + input_approx, + backend=ctx.config.matmul_backend, + input_config=wgrad_grad_config, + other_config=wgrad_activation_config, + out_dtype=ctx.config.output_dtype, + ).unsqueeze(0) + + grad_bias = ( + grad_output.sum(0) if bias is not None and ctx.needs_input_grad[3] else None + ) + + return ( + None, + grad_input, + grad_weight, + grad_bias, + ) + +@QuantizedModule.register(nn.Linear) +class FourOverSixLinear(nn.Linear): + """ + Drop-in replacement for `nn.Linear` that quantizes weights, activations, and + gradients. + """ + def __init__( + self, + module: nn.Linear, + config: ModuleQuantizationConfig, + ) -> None: + """ + Initialize the FourOverSixLinear layer. + Args: + module (nn.Linear): The high-precision module that this quantized layer will + replace. + config (ModuleQuantizationConfig): The quantization configuration to use for + the layer. + """ + super().__init__( + module.in_features, + module.out_features, + module.bias is not None, + module.weight.device, + module.weight.dtype, + ) + self.weight = module.weight + self.bias = module.bias + self.config = config + if not self.config.keep_master_weights: + self.register_buffer( + "quantized_weight_values", + nn.Parameter( + torch.zeros( + self.out_features, + self.in_features // 2, + dtype=torch.uint8, + ), + requires_grad=False, + ), + ) + self.register_buffer( + "quantized_weight_scale_factors", + nn.Parameter( + torch.zeros( + self.out_features + * self.in_features + // self.config.dtype.block_size(), + dtype=self.config.dtype.scale_dtype(), + ), + requires_grad=False, + ), + ) + self.register_buffer( + "quantized_weight_amax", + nn.Parameter(torch.zeros(1, dtype=torch.float32), requires_grad=False), + ) + self.register_buffer( + "quantized_weight_metadata", + nn.Parameter( + torch.zeros(2 + 2, dtype=torch.int32), + requires_grad=False, + ), + ) + + @property + def parameters_to_quantize(self) -> tuple[str, ...]: + """Return high precision parameters to be quantized and deleted.""" + return ("weight",) + + def get_element_size(self, parameter_name: str) -> float: + """Get the size of a single element, in bytes, for a parameter.""" + # quantized_weight_values is packed, so there are 4 bits, or 0.5 bytes, per + # element. Once quantized, weight will have (8+1)/16 bytes per element (one + # block of 16 values is 8 bytes of values + 1 byte of scale factors). + return {"quantized_weight_values": 0.5, "weight": 9 / 16}.get( + parameter_name, + getattr(self, parameter_name).element_size(), + ) + + def get_quantized_parameters( + self, + parameter_name: str, + parameter: torch.Tensor, + ) -> dict[str, Any]: + """Get the quantized parameters for the layer.""" + if parameter_name == "weight": + config = QuantizationConfig( + backend=self.config.quantize_backend, + block_scale_2d=self.config.weight_scale_2d, + dtype=self.config.dtype, + scale_rule=self.config.weight_scale_rule, + ) + quantized_weight = quantize_to_fp4(parameter, config) + return { + "quantized_weight_values": quantized_weight.values, + "quantized_weight_scale_factors": quantized_weight.scale_factors, + "quantized_weight_amax": quantized_weight.amax, + "quantized_weight_metadata": torch.tensor( + [ + quantized_weight.original_shape[0], + quantized_weight.original_shape[1], + quantized_weight.padded_shape[0], + quantized_weight.padded_shape[1], + ], + ), + } + msg = f"Unsupported high-preciison parameter: {parameter_name}" + raise ValueError(msg) + + def quantized_weight(self) -> QuantizedTensor: + """ + Prepare this layer for post-training quantization by quantizing the weight, + storing the quantized weight, and deleting the original weight. This should not + be done if the layer is used for training, as training requires storage of the + high-precision weight. + """ + if not hasattr(self, "_quantized_weight"): + if self.config.keep_master_weights: + return self.weight + original_shape = tuple(self.quantized_weight_metadata.data[:2].tolist()) + padded_shape = tuple(self.quantized_weight_metadata.data[2:].tolist()) + self._quantized_weight = QuantizedTensor( + self.quantized_weight_values.data, + self.quantized_weight_scale_factors.data, + self.quantized_weight_amax.data, + self.config.dtype, + original_shape, + self.config.weight_scale_rule, + padded_shape, + ) + return self._quantized_weight + + def forward(self, input: torch.Tensor) -> torch.Tensor: + """Forward pass for the FP4 linear layer.""" + return FourOverSixLinearFunction.apply( + self.config, + input, + self.quantized_weight(), + self.bias, + ) \ No newline at end of file diff --git a/fouroversix/src/fouroversix/model/modules/qwen.py b/fouroversix/src/fouroversix/model/modules/qwen.py new file mode 100644 index 0000000..95d0cf6 --- /dev/null +++ b/fouroversix/src/fouroversix/model/modules/qwen.py @@ -0,0 +1,326 @@ +import warnings +from typing import Any + +import torch +from fouroversix.matmul import fp4_matmul +from fouroversix.model.config import ModuleQuantizationConfig +from fouroversix.model.quantize import QuantizedModule +from fouroversix.quantize import QuantizationConfig, QuantizedTensor, quantize_to_fp4 +from torch import nn + +try: + from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeExperts +except ImportError: + warnings.warn( + "Qwen3_5MoeExperts not found, please update transformers to the latest " + "version to quantize Qwen3.5 models.", + stacklevel=2, + ) + + Qwen3_5MoeExperts = None + + +@QuantizedModule.register(Qwen3_5MoeExperts) +class FourOverSixQwenExperts(nn.Module): + """ + Drop-in replacement for the Qwen3_5MoeExperts layer that + uses FP4 quantization. + """ + + def __init__( + self, + module: Qwen3_5MoeExperts, + config: ModuleQuantizationConfig, + ) -> None: + """ + Initialize the FourOverSixQwenExperts layer. + + Args: + module (GptOssMLP): The high-precision module that this quantized layer will + replace. + config (ModuleQuantizationConfig): The quantization configuration to use for + the layer. + + """ + super().__init__() + + self.num_experts = module.num_experts + self.intermediate_dim = module.intermediate_dim + self.hidden_dim = module.hidden_dim + + self.down_proj = module.down_proj + self.gate_up_proj = module.gate_up_proj + + self.device = self.down_proj.device + self.config = config + + self.act_fn = module.act_fn + + if not self.config.keep_master_weights: + + self.register_buffer( + "quantized_down_proj_values", + nn.Parameter( + torch.zeros( + self.num_experts, + self.hidden_dim, + self.intermediate_dim // 2, + dtype=torch.uint8, + ), + requires_grad=False, + ), + ) + + self.register_buffer( + "quantized_gate_up_proj_values", + nn.Parameter( + torch.zeros( + self.num_experts, + self.intermediate_dim * 2, + self.hidden_dim // 2, + dtype=torch.uint8, + ), + requires_grad=False, + ), + ) + + self.register_buffer( + "quantized_down_proj_scale_factors", + nn.Parameter( + torch.zeros( + self.num_experts, + self.hidden_dim + * self.intermediate_dim + // self.config.dtype.block_size(), + dtype=self.config.dtype.scale_dtype(), + ), + requires_grad=False, + ), + ) + self.register_buffer( + "quantized_gate_up_proj_scale_factors", + nn.Parameter( + torch.zeros( + self.num_experts, + self.hidden_dim + * (self.intermediate_dim * 2) + // self.config.dtype.block_size(), + dtype=self.config.dtype.scale_dtype(), + ), + requires_grad=False, + ), + ) + + self.register_buffer( + "quantized_down_proj_amax", + nn.Parameter( + torch.zeros(self.num_experts, 1, dtype=torch.float32), + requires_grad=False, + ), + ) + self.register_buffer( + "quantized_gate_up_proj_amax", + nn.Parameter( + torch.zeros(self.num_experts, 1, dtype=torch.float32), + requires_grad=False, + ), + ) + + self.register_buffer( + "quantized_down_proj_metadata", + nn.Parameter( + torch.zeros(self.num_experts, 4, dtype=torch.int32), + requires_grad=False, + ), + ) + self.register_buffer( + "quantized_gate_up_proj_metadata", + nn.Parameter( + torch.zeros(self.num_experts, 4, dtype=torch.int32), + requires_grad=False, + ), + ) + + @property + def parameters_to_quantize(self) -> tuple[str, ...]: + """Return high precision parameters to be quantized and deleted.""" + return ("down_proj", "gate_up_proj") + + def get_element_size(self, parameter_name: str) -> float: + """Get the size of a single element, in bytes, for a parameter.""" + + return { + "quantized_down_proj_values": 0.5, + "quantized_gate_up_proj_values": 0.5, + "down_proj": 9 / 16, + "gate_up_proj": 9 / 16, + }.get( + parameter_name, + getattr(self, parameter_name).element_size(), + ) + + def get_quantized_parameters( + self, + parameter_name: str, + parameter: torch.Tensor, + ) -> dict[str, Any]: + """ + Prepare this layer for post-training quantization by quantizing the weight, + storing the quantized weight, and deleting the original weight. This should not + be done if the layer is used for training, as training requires storage of the + high-precision weight. + """ + + if "bias" in parameter_name: + return {parameter_name: parameter} + + weight_config = QuantizationConfig( + backend=self.config.quantize_backend, + dtype=self.config.dtype, + scale_rule=self.config.weight_scale_rule, + ) + + quantized_proj = [] + for e in range(parameter.shape[0]): + q = quantize_to_fp4(parameter[e], weight_config) + quantized_proj.append(q) + + if "down" in parameter_name: + prefix = "down" + elif "gate_up" in parameter_name: + prefix = "gate_up" + + return { + f"quantized_{prefix}_proj_values": torch.stack( + [tensor.values for tensor in quantized_proj], + dim=0, + ), + f"quantized_{prefix}_proj_scale_factors": torch.stack( + [tensor.scale_factors for tensor in quantized_proj], + dim=0, + ), + f"quantized_{prefix}_proj_amax": torch.stack( + [tensor.amax for tensor in quantized_proj], + dim=0, + ), + f"quantized_{prefix}_proj_metadata": torch.stack( + [ + torch.tensor( + [ + tensor.original_shape[0], + tensor.original_shape[1], + tensor.padded_shape[0], + tensor.padded_shape[1], + ], + ) + for tensor in quantized_proj + ], + ), + } + + def forward( + self, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, + ) -> torch.Tensor: + """Forward pass for the FP4 experts layer.""" + + down_proj, gate_up_proj = self.quantized_weights() + + final_hidden_states = torch.zeros_like(hidden_states) + with torch.no_grad(): + expert_mask = torch.nn.functional.one_hot( + top_k_index, + num_classes=self.num_experts, + ) + expert_mask = expert_mask.permute(2, 1, 0) + expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() + + for [expert_idx] in expert_hit: + if expert_idx == self.num_experts: + continue + + fprop_activation_config = QuantizationConfig( + backend=self.config.quantize_backend, + dtype=self.config.dtype, + scale_rule=self.config.activation_scale_rule, + ) + top_k_pos, token_idx = torch.where(expert_mask[expert_idx]) + current_state = hidden_states[token_idx] + gate, up = fp4_matmul( + current_state, + gate_up_proj[expert_idx], + input_config=fprop_activation_config, + out_dtype=self.config.output_dtype, + ).chunk(2, dim=-1) + current_hidden_states = self.act_fn(gate) * up + current_hidden_states = fp4_matmul( + current_hidden_states, + down_proj[expert_idx], + input_config=fprop_activation_config, + out_dtype=self.config.output_dtype, + ) + current_hidden_states = ( + current_hidden_states * top_k_weights[token_idx, top_k_pos, None] + ) + final_hidden_states.index_add_( + 0, + token_idx, + current_hidden_states.to(final_hidden_states.dtype), + ) + + return final_hidden_states + + def quantized_weights(self) -> tuple[list[QuantizedTensor], list[QuantizedTensor]]: + """Return quantized parameters as QuantizedTensor.""" + + if not hasattr(self, "_quantized_weights"): + weight_config = self.config.get_weight_config() + if self.config.keep_master_weights: + down = [ + quantize_to_fp4(self.down_proj[e], weight_config) + for e in range(self.num_experts) + ] + gate_up = [ + quantize_to_fp4(self.gate_up_proj[e], weight_config) + for e in range(self.num_experts) + ] + return (down, gate_up) + + down = [] + gate_up = [] + for e in range(self.num_experts): + down.append( + QuantizedTensor( + values=self.quantized_down_proj_values.data[e], + scale_factors=self.quantized_down_proj_scale_factors.data[e], + amax=self.quantized_down_proj_amax.data[e], + dtype=self.config.dtype, + original_shape=tuple( + self.quantized_down_proj_metadata.data[e, :2].tolist(), + ), + scale_rule=self.config.weight_scale_rule, + padded_shape=tuple( + self.quantized_down_proj_metadata.data[e, 2:].tolist(), + ), + ), + ) + gate_up.append( + QuantizedTensor( + values=self.quantized_gate_up_proj_values.data[e], + scale_factors=self.quantized_gate_up_proj_scale_factors.data[e], + amax=self.quantized_gate_up_proj_amax.data[e], + dtype=self.config.dtype, + original_shape=tuple( + self.quantized_gate_up_proj_metadata.data[e, :2].tolist(), + ), + scale_rule=self.config.weight_scale_rule, + padded_shape=tuple( + self.quantized_gate_up_proj_metadata.data[e, 2:].tolist(), + ), + ), + ) + self._quantized_weights = (down, gate_up) + + return self._quantized_weights diff --git a/fouroversix/src/fouroversix/model/quantize.py b/fouroversix/src/fouroversix/model/quantize.py new file mode 100644 index 0000000..444bda7 --- /dev/null +++ b/fouroversix/src/fouroversix/model/quantize.py @@ -0,0 +1,127 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, ClassVar + +import torch.nn as nn + +if TYPE_CHECKING: + from collections.abc import Callable + + from .config import ModelQuantizationConfig + + +class QuantizedModule: + """Base class for all quantized modules.""" + + _registry: ClassVar[dict[type[nn.Module], type[nn.Module]]] = {} + _should_replace_existing_modules_in_model: ClassVar[dict[type[nn.Module], bool]] = ( + {} + ) + + @classmethod + def is_quantized_module_type(cls, module_type: type[nn.Module]) -> bool: + """Return True if the given module type is a quantized module.""" + return module_type in cls._registry.values() + + @classmethod + def get_cls( + cls, + high_precision_cls: type[nn.Module], + ) -> type[nn.Module] | None: + """Get the quantized module for a given high-precision module.""" + return cls._registry.get(high_precision_cls) + + @classmethod + def should_replace_existing_modules_in_model( + cls, + module_type: type[nn.Module], + ) -> bool: + """Determine whether module should be replaced.""" + return cls._should_replace_existing_modules_in_model.get(module_type, False) + + @classmethod + def register( + cls, + high_precision_cls: type[nn.Module], + *, + replace_existing_modules_in_registry: bool = False, + replace_existing_modules_in_model: bool = True, + ) -> Callable[[type[nn.Module]], type[nn.Module]]: + """ + Register a new type of quantized module. + + Args: + high_precision_cls: (`type[nn.Module]`): The high precision module to be + mapped to a fouroversix quantized module. + replace_existing_modules_in_registry (bool): determines whether we should + replace the existing module in the registry. + replace_existing_modules_in_model (bool): determines whether we should + replace the existing module in the model including the weights. + + """ + + if ( + high_precision_cls in cls._registry + and not replace_existing_modules_in_registry + ): + msg = f"High-precision module {high_precision_cls} is already registered." + raise ValueError(msg) + + modules_to_delete = [] + + for module_cls in cls._registry: + if high_precision_cls is not None and issubclass( + high_precision_cls, + module_cls, + ): + if replace_existing_modules_in_registry: + modules_to_delete.append(module_cls) + else: + msg = ( + f"High-precision module {high_precision_cls} is a subclass of " + f"{module_cls}, which is already registered." + ) + raise TypeError(msg) + + for module_cls in modules_to_delete: + del cls._registry[module_cls] + + def inner_wrapper( + wrapped_cls: type[nn.Module], + ) -> type[nn.Module]: + cls._registry[high_precision_cls] = wrapped_cls + cls._should_replace_existing_modules_in_model[high_precision_cls] = ( + replace_existing_modules_in_model + ) + return wrapped_cls + + return inner_wrapper + + +def quantize_model( + model: nn.Module, + config: ModelQuantizationConfig, + **kwargs: dict[str, Any], +) -> None: + for module_name, module in model.named_modules(): + if ( + module_name == "" + or module_name in config.modules_to_not_convert + or not isinstance(module, nn.Module) + ): + continue + + module_cls = QuantizedModule.get_cls(type(module)) + should_replace = QuantizedModule.should_replace_existing_modules_in_model( + type(module), + ) + + if module_cls is None or not should_replace: + continue + + quantized_module = module_cls( + module, + config.get_module_config(module_name), + **kwargs, + ) + model.set_submodule(module_name, quantized_module) diff --git a/fouroversix/src/fouroversix/quantize/__init__.py b/fouroversix/src/fouroversix/quantize/__init__.py new file mode 100644 index 0000000..afd3ad6 --- /dev/null +++ b/fouroversix/src/fouroversix/quantize/__init__.py @@ -0,0 +1,6 @@ +from .config import QuantizationConfig +from .frontend import quantize_to_fp4 +from .quantized_tensor import QuantizedTensor +from .utils import get_rht_matrix + +__all__ = ["QuantizationConfig", "QuantizedTensor", "get_rht_matrix", "quantize_to_fp4"] diff --git a/fouroversix/src/fouroversix/quantize/backend.py b/fouroversix/src/fouroversix/quantize/backend.py new file mode 100644 index 0000000..e47cd3d --- /dev/null +++ b/fouroversix/src/fouroversix/quantize/backend.py @@ -0,0 +1,69 @@ +from abc import ABC, abstractmethod + +import torch +from fouroversix.utils import DataType, ScaleRule + +from .config import QuantizationConfig +from .quantized_tensor import QuantizedTensor + + +class QuantizeBackendBase(ABC): + """Base class for all quantization backends.""" + + @classmethod + @abstractmethod + def is_available(cls) -> bool: + """Return True if the backend is available on the current machine.""" + msg = "Subclasses must implement this method" + raise NotImplementedError(msg) + + @classmethod + @abstractmethod + def is_supported(cls, x: torch.Tensor, config: QuantizationConfig) -> bool: + """ + Return True if the backend supports the given input and quantization + configuration. + """ + + if not cls.is_available(): + return False + + if x.ndim != 2: # noqa: PLR2004 + return False + + if config.dtype not in {DataType.mxfp4, DataType.nvfp4}: + return False + + if config.dtype == DataType.mxfp4 and config.scale_rule not in { + ScaleRule.static_6, + ScaleRule.static_4, + }: + msg = ( + "MXFP4 quantization only supports the `static_6` and `static_4` scale " + "rules" + ) + raise ValueError(msg) + + return True + + @classmethod + @abstractmethod + def quantize_to_fp4( + cls, + x: torch.Tensor, + config: QuantizationConfig, + ) -> QuantizedTensor: + """ + Quantize a tensor to FP4 using the backend. + + Args: + x (torch.Tensor): The input tensor to quantize. + config (QuantizationConfig): The quantization configuration. + + Returns: + The quantized tensor. + + """ + + msg = "Subclasses must implement this method" + raise NotImplementedError(msg) diff --git a/fouroversix/src/fouroversix/quantize/config.py b/fouroversix/src/fouroversix/quantize/config.py new file mode 100644 index 0000000..cc5d1ee --- /dev/null +++ b/fouroversix/src/fouroversix/quantize/config.py @@ -0,0 +1,51 @@ +from dataclasses import dataclass + +from fouroversix.utils import DataType, QuantizeBackend, RoundStyle, ScaleRule + + +@dataclass +class QuantizationConfig: + """ + Configuration to use when quantizing a tensor. + + Args: + backend (QuantizeBackend): The backend to use for quantization. If no backend is + provided, a backend will be selected automatically based on the available + GPU and the specified options. + block_scale_2d (bool): If True, scale factors will be computed across 16x16 + chunks of the input rather than 1x16 chunks. This is useful to apply to the + weight matrix during training, so that W and W.T will be equivalent after + quantization. + dtype (DataType): The data type to quantize to. + rht (bool): If True, the random Hadamard transform will be applied to the input + prior to quantization. + round_style (RoundStyle): The rounding style to apply during quantization. + scale_rule (ScaleRule): The scaling rule to use during quantization. + transpose (bool): If True, the output will be a quantized version of the + transposed input. This may be helpful for certain operations during training + as `fp4_matmul` requires that both tensors are provided in row-major format. + + """ + + backend: QuantizeBackend | None = None + block_scale_2d: bool = False + dtype: DataType = DataType.nvfp4 + rht: bool = False + round_style: RoundStyle = RoundStyle.nearest + scale_rule: ScaleRule = ScaleRule.mse + transpose: bool = False + + def __post_init__(self) -> None: + """Convert string values to enums.""" + + if isinstance(self.backend, str): + self.backend = QuantizeBackend(self.backend) + + if isinstance(self.dtype, str): + self.dtype = DataType(self.dtype) + + if isinstance(self.round_style, str): + self.round_style = RoundStyle(self.round_style) + + if isinstance(self.scale_rule, str): + self.scale_rule = ScaleRule(self.scale_rule) diff --git a/fouroversix/src/fouroversix/quantize/cuda/__init__.py b/fouroversix/src/fouroversix/quantize/cuda/__init__.py new file mode 100644 index 0000000..deb5bf3 --- /dev/null +++ b/fouroversix/src/fouroversix/quantize/cuda/__init__.py @@ -0,0 +1,3 @@ +from .backend import CUDAQuantizeBackend + +__all__ = ["CUDAQuantizeBackend"] diff --git a/fouroversix/src/fouroversix/quantize/cuda/backend.py b/fouroversix/src/fouroversix/quantize/cuda/backend.py new file mode 100644 index 0000000..89bb185 --- /dev/null +++ b/fouroversix/src/fouroversix/quantize/cuda/backend.py @@ -0,0 +1,91 @@ +import functools + +import torch +from fouroversix.quantize.backend import QuantizeBackendBase +from fouroversix.quantize.config import QuantizationConfig +from fouroversix.quantize.quantized_tensor import QuantizedTensor +from fouroversix.utils import BLACKWELL_SM_IDS, DataType, RoundStyle + + +class CUDAQuantizeBackend(QuantizeBackendBase): + """ + The CUDA quantization backend. Supports basic quantization options (no 2D block + scaling, no stochastic rounding, no random Hadamard transform). As a result, it can + be used for inference, but not training. Requires a Blackwell GPU. + """ + + @classmethod + @functools.lru_cache + def is_available(cls) -> bool: + """Return True if the CUDA backend is available on the current machine.""" + + if ( + not torch.cuda.is_available() + or torch.cuda.get_device_capability()[0] not in BLACKWELL_SM_IDS + ): + return False + + try: + import fouroversix._C # noqa: F401 + except ModuleNotFoundError: + return False + + return True + + @classmethod + def is_supported(cls, x: torch.Tensor, config: QuantizationConfig) -> bool: + """ + Return True if the CUDA backend supports the given input and quantization + configuration. + """ + + if not super().is_supported(x, config): + return False + + return ( + x.device.type == "cuda" + and x.dtype in {torch.float16, torch.bfloat16} + and config.round_style == RoundStyle.nearest + and config.dtype == DataType.nvfp4 + and not config.transpose + ) + + @classmethod + def quantize_to_fp4( + cls, + x: torch.Tensor, + config: QuantizationConfig, + ) -> QuantizedTensor: + """ + Quantize a tensor to FP4 using the CUDA backend. + + Args: + x (torch.Tensor): The input tensor to quantize. + config (QuantizationConfig): The quantization configuration. + + Returns: + The quantized tensor. + + """ + + from .ops import quantize_to_fp4 + + values, scale_factors, amax = quantize_to_fp4( + x, + config.dtype == DataType.nvfp4, + config.round_style == RoundStyle.nearest, + config.rht, + config.block_scale_2d, + config.transpose, + config.scale_rule.cuda_id(), + 0, + ) + + return QuantizedTensor( + values, + scale_factors, + amax, + config.dtype, + (x.shape[1], x.shape[0]) if config.transpose else x.shape, + config.scale_rule, + ) \ No newline at end of file diff --git a/fouroversix/src/fouroversix/quantize/cuda/ops.py b/fouroversix/src/fouroversix/quantize/cuda/ops.py new file mode 100644 index 0000000..d7a0844 --- /dev/null +++ b/fouroversix/src/fouroversix/quantize/cuda/ops.py @@ -0,0 +1,46 @@ +import torch +import fouroversix._C # noqa: F401 + +def quantize_to_fp4( + x: torch.Tensor, + is_nvfp4: bool, # noqa: FBT001 + is_rtn: bool, # noqa: FBT001 + is_rht: bool, # noqa: FBT001 + is_2d: bool, # noqa: FBT001 + is_transpose: bool, # noqa: FBT001 + selection_rule: int, + rbits: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + + return torch.ops.fouroversix.quantize_to_fp4.default( + x, + is_nvfp4, + is_rtn, + is_rht, + is_2d, + is_transpose, + selection_rule, + rbits, + ) + +@torch.library.register_fake("fouroversix::quantize_to_fp4") +def _( + x: torch.Tensor, + is_nvfp4: bool, # noqa: FBT001 + is_rtn: bool, # noqa: ARG001, FBT001 + is_rht: bool, # noqa: ARG001, FBT001 + is_2d: bool, # noqa: ARG001, FBT001 + is_transpose: bool, # noqa: ARG001, FBT001 + selection_rule: int, # noqa: ARG001 + rbits: int, # noqa: ARG001 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + + return ( + torch.empty(x.shape[0], x.shape[1] // 2, dtype=torch.uint8, device=x.device), + torch.empty( + x.shape[0] * x.shape[1] // (16 if is_nvfp4 else 32), + dtype=torch.float8_e4m3fn if is_nvfp4 else torch.float8_e8m0fnu, + device=x.device, + ), + torch.empty(1, dtype=torch.float32, device=x.device), + ) \ No newline at end of file diff --git a/fouroversix/src/fouroversix/quantize/frontend.py b/fouroversix/src/fouroversix/quantize/frontend.py new file mode 100644 index 0000000..f5b9f00 --- /dev/null +++ b/fouroversix/src/fouroversix/quantize/frontend.py @@ -0,0 +1,122 @@ +import torch +from fouroversix.utils import QuantizeBackend + +from .config import QuantizationConfig +from .cuda import CUDAQuantizeBackend +from .pytorch import PyTorchQuantizeBackend +from .quantized_tensor import QuantizedTensor +from .transformer_engine import TransformerEngineQuantizeBackend +from .triton import TritonQuantizeBackend + +AVAILABLE_BACKENDS = { + QuantizeBackend.cuda: CUDAQuantizeBackend, + QuantizeBackend.transformer_engine: TransformerEngineQuantizeBackend, + QuantizeBackend.triton: TritonQuantizeBackend, + QuantizeBackend.pytorch: PyTorchQuantizeBackend, +} + + +def quantize_to_fp4( + x: torch.Tensor, + config: QuantizationConfig | None = None, +) -> QuantizedTensor: + """ + Quantize a tensor to FP4. + + ## Sample Code + + ### With Four Over Six + + ```python + x = torch.tensor(1024, 1024, dtype=torch.bfloat16, device="cuda") + x_quantized = quantize_to_fp4(x) + ``` + + ### Without Four Over Six + + ```python + x = torch.tensor(1024, 1024, dtype=torch.bfloat16, device="cuda") + config = QuantizationConfig(scale_rule="static_6") + x_quantized = quantize_to_fp4(x, config) + ``` + + ### With Stochastic Rounding + + ```python + x = torch.tensor(1024, 1024, dtype=torch.bfloat16, device="cuda") + config = QuantizationConfig(round_style="stochastic") + x_quantized = quantize_to_fp4(x, config) + ``` + + ### With the Random Hadamard Transform + + ```python + from fouroversix.quantize import get_rht_matrix + + x = torch.tensor(1024, 1024, dtype=torch.bfloat16, device="cuda") + config = QuantizationConfig(rht=True) + x_quantized = quantize_to_fp4(x, config) + ``` + + ## Backends + + We provide three different implementations of FP4 quantization: + + - **CUDA**: A fast implementation written in CUDA which currently only supports + basic quantization options (no 2D block scaling, no stochastic rounding, no + random Hadamard transform). Can be used for inference, but not training. + Requires a Blackwell GPU. + - **Triton**: A slightly slower implementation written in Triton which supports all + operations needed for training. Also requires a Blackwell GPU. + - **PyTorch**: A slow implementation written in PyTorch which supports all + operations and can be run on any GPU. + + If `quantize_to_fp4` is called with `backend=None`, a backend will be selected + automatically based on the following rules: + + - If there is no GPU available, or if the available GPU is not a Blackwell GPU, + select PyTorch. + - If any quantization options are set other than `scale_rule`, select Triton. + - However, if the available GPU is SM120 (i.e. RTX 5090, RTX 6000) and + `round_style` is set to `RoundStyle.stochastic`, select PyTorch as + stochastic rounding does not have hardware support on SM120 GPUs. + - Otherwise, select CUDA. + + ## Parameters + + Args: + x (torch.Tensor): The input tensor to quantize. + config (QuantizationConfig): The quantization configuration to use. If no + configuration is provided, a default configuration will be used (NVFP4, + 1D block scaling, round-to-nearest, and 4/6 with the MSE selection rule). + + Returns: + The quantized tensor. + + """ + + if config is None: + config = QuantizationConfig() + + selected_backend = config.backend + + if selected_backend is None: + for backend in [ + QuantizeBackend.cuda, + QuantizeBackend.triton, + QuantizeBackend.pytorch, + ]: + if AVAILABLE_BACKENDS[backend] is not None and AVAILABLE_BACKENDS[ + backend + ].is_supported(x, config): + selected_backend = backend + break + else: + msg = "No backend found that supports the given parameters" + raise ValueError(msg) + + elif not AVAILABLE_BACKENDS[selected_backend].is_supported(x, config): + msg = f"Backend {selected_backend} does not support the given parameters" + raise ValueError(msg) + + return AVAILABLE_BACKENDS[selected_backend].quantize_to_fp4(x, config) diff --git a/fouroversix/src/fouroversix/quantize/pytorch/__init__.py b/fouroversix/src/fouroversix/quantize/pytorch/__init__.py new file mode 100644 index 0000000..b0d51dc --- /dev/null +++ b/fouroversix/src/fouroversix/quantize/pytorch/__init__.py @@ -0,0 +1,3 @@ +from .backend import PyTorchQuantizeBackend + +__all__ = ["PyTorchQuantizeBackend"] diff --git a/fouroversix/src/fouroversix/quantize/pytorch/backend.py b/fouroversix/src/fouroversix/quantize/pytorch/backend.py new file mode 100644 index 0000000..0ced49e --- /dev/null +++ b/fouroversix/src/fouroversix/quantize/pytorch/backend.py @@ -0,0 +1,112 @@ +import torch +import torch.nn.functional as F +from fouroversix.quantize.backend import QuantizeBackendBase +from fouroversix.quantize.config import QuantizationConfig +from fouroversix.quantize.quantized_tensor import QuantizedTensor +from fouroversix.quantize.utils import get_rht_matrix + +from .reference import quantize_to_fp4 + + +class PyTorchQuantizeBackend(QuantizeBackendBase): + """ + The PyTorch quantization backend. Supports all quantization options, and can be run + on non-Blackwell GPUs, but is slow. Should be used primarily as a reference. + """ + + @classmethod + def is_available(cls) -> bool: + """Return True if the PyTorch backend is available on the current machine.""" + return True + + @classmethod + def is_supported( + cls, + x: torch.Tensor, # noqa: ARG003 + config: QuantizationConfig, # noqa: ARG003 + ) -> bool: + """ + Return True if the PyTorch backend supports the given input and quantization + configuration. + """ + + return True + + @classmethod + def quantize_to_fp4( + cls, + x: torch.Tensor, + config: QuantizationConfig, + ) -> QuantizedTensor: + """ + Quantize a tensor to FP4 using the PyTorch backend. + + Args: + x (torch.Tensor): The input tensor to quantize. + config (QuantizationConfig): The quantization configuration. + + Returns: + The quantized tensor. + + """ + + input_shape = (x.shape[1], x.shape[0]) if config.transpose else x.shape + + rows_div = 128 + cols_div = 4 * config.dtype.block_size() + + if input_shape[0] % rows_div != 0 or input_shape[1] % cols_div != 0: + x = F.pad( + x, + ( + 0, + ( + cols_div - (input_shape[1] % cols_div) + if input_shape[1] % cols_div > 0 + else 0 + ), + 0, + ( + rows_div - (input_shape[0] % rows_div) + if input_shape[0] % rows_div > 0 + else 0 + ), + ), + ) + + if x.device.type == "meta": + values = torch.zeros( + input_shape[0], + input_shape[1] // 2, + device=x.device, + dtype=torch.uint8, + ) + scale_factors = torch.zeros( + input_shape[0] * input_shape[1] // config.dtype.block_size(), + device=x.device, + dtype=( + torch.uint8 + if config.dtype.scale_dtype() == torch.float8_e8m0fnu + else config.dtype.scale_dtype() + ), + ) + amax = torch.zeros(1, device=x.device, dtype=torch.float32) + else: + values, scale_factors, amax = quantize_to_fp4( + x, + had=get_rht_matrix() if config.rht else None, + fp4_format=config.dtype, + round_style=config.round_style, + scale_rule=config.scale_rule, + block_scale_2d=config.block_scale_2d, + transpose=config.transpose, + ) + + return QuantizedTensor( + values, + scale_factors, + amax, + config.dtype, + input_shape, + config.scale_rule, + ) diff --git a/fouroversix/src/fouroversix/quantize/pytorch/reference.py b/fouroversix/src/fouroversix/quantize/pytorch/reference.py new file mode 100644 index 0000000..66852ca --- /dev/null +++ b/fouroversix/src/fouroversix/quantize/pytorch/reference.py @@ -0,0 +1,368 @@ +from __future__ import annotations + +import torch +from fouroversix.quantize.utils import to_blocked +from fouroversix.utils import DataType, RoundStyle, ScaleRule + +E2M1_MAX_VALUE = 6 +E2M1_MAX_FOUR = 4 +E4M3_MAX_VALUE = 448 +E4M3_MAX_FOUROVERSIX = 256 + + +def fake_quantize_to_e2m1( + x: torch.Tensor, + *, + round_style: RoundStyle = RoundStyle.nearest, +) -> torch.Tensor: + if round_style == RoundStyle.nearest: + step1 = torch.round(2 * x.abs()) / 2 + step2 = torch.round(x.abs()) + step3 = 2 * torch.round(x.abs() / 2) + elif round_style == RoundStyle.stochastic: + rbits = torch.rand_like(x.abs()) - 0.5 + step1 = torch.round(2 * x.abs() + rbits) / 2 + step2 = torch.round(x.abs() + rbits) + step3 = 2 * torch.round(x.abs() / 2 + rbits) + step3[step3 > E2M1_MAX_VALUE] = E2M1_MAX_VALUE + + mask1 = x.abs() < 2 # noqa: PLR2004 + mask2 = x.abs() < 4 # noqa: PLR2004 + + return x.sign() * ( + step1 * mask1 + step2 * (~mask1) * mask2 + step3 * (~mask1) * (~mask2) + ) + + +def quantize_bf16_to_unpacked_fp4(x: torch.Tensor) -> torch.Tensor: + assert x.dtype == torch.bfloat16 + + bx = x.view(torch.int16) + s = (bx >> 15) & 0x1 + e = (bx >> 7) & 0xFF + m = bx & 0x7F + is_zero = (e == 0) & (m == 0) + + # Default mantissa bit (for 1.5, 3.0, 6.0) + m = (m >> 6) & 1 + is_half = (e == 126) & (m == 0) # noqa: PLR2004 + m = torch.where(is_half, torch.tensor(1, dtype=torch.int16, device=x.device), m) + + # Exponent mapping + # exp=126 -> E=0 (subnormals) + # exp=127 -> E=1 + # exp=128 -> E=2 + # exp=129 -> E=3 + e = e - 126 + e = torch.where(is_zero, torch.tensor(0, dtype=torch.int16, device=x.device), e) + + # Zero always M=0 + m = torch.where(is_zero, torch.tensor(0, dtype=torch.int16, device=x.device), m) + + code = (s << 3) | (e << 1) | m + return code.to(torch.uint8) + + +def pack_unpacked_fp4(x: torch.Tensor) -> torch.Tensor: + assert x.dtype == torch.uint8 + + dim = 1 + size_along_dim = x.size(dim) + new_size_along_dim = (size_along_dim + 1) // 2 + + # If the size is odd, we pad the data along dim with zeros at the end + if size_along_dim % 2 != 0: + pad_sizes = [0] * (2 * x.ndim) + pad_index = (x.ndim - dim - 1) * 2 + 1 + pad_sizes[pad_index] = 1 + x = torch.nn.functional.pad(x, pad_sizes, mode="constant", value=0) + + new_shape = list(x.shape) + new_shape[dim] = new_size_along_dim + new_shape.insert(dim + 1, 2) # packed dimension of length 2 + x = x.reshape(*new_shape) + + low = x.select(dim + 1, 0) + high = x.select(dim + 1, 1) + return (high << 4) | low + + +def quantize_to_mxfp4( + x_scale_blocks: torch.Tensor, + *, + scale_rule: ScaleRule = ScaleRule.mse, +) -> tuple[torch.Tensor, torch.Tensor]: + assert scale_rule in {ScaleRule.static_6, ScaleRule.static_4} + + x_scales_hp = ( + x_scale_blocks.abs().max(axis=-1).values / scale_rule.max_allowed_e2m1_value() + ) + + x_scales_e8m0_u32 = x_scales_hp.view(torch.int32) + + # Use the 8-bit exponent as the scale factor + x_scales_e8m0 = ((x_scales_e8m0_u32 >> 23) & 0xFF).to(torch.uint8) + + # Add one in order to round up + x_scales = torch.where( + (x_scales_e8m0_u32 & 0x7FFFFF) == 0, + x_scales_e8m0, + x_scales_e8m0 + 1, + ) + + # Convert the rounded-up scale factor back to a 32-bit float + x_scales_hp = (x_scales.to(torch.int32) << 23).view(torch.float32) + + x_block_scaled = x_scale_blocks / x_scales_hp.unsqueeze(1) + + return x_block_scaled, x_scales.view(torch.float8_e8m0fnu) + + +def quantize_to_nvfp4( + x_scale_blocks: torch.Tensor, + x_amax: torch.Tensor, + *, + scale_rule: ScaleRule, + scale_expansion_factor: float | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + if x_amax == 0: + x_scales_hp = torch.zeros( + *x_scale_blocks.shape[:-1], + dtype=x_amax.dtype, + device=x_amax.device, + ) + else: + encode_scale = ( + torch.tensor( + scale_rule.max_allowed_e2m1_value() + * scale_rule.max_allowed_e4m3_value(), + dtype=x_amax.dtype, + device=x_amax.device, + ) + / x_amax + ) + x_scales_hp = ( + x_scale_blocks.abs().max(axis=-1).values + / torch.tensor( + scale_rule.max_allowed_e2m1_value(), + dtype=x_amax.dtype, + device=x_amax.device, + ) + * encode_scale + ) + + if scale_expansion_factor is not None: + x_scales_hp = x_scales_hp * scale_expansion_factor + + x_scales = x_scales_hp.to(torch.float8_e4m3fn) + + decode_scale = 1 / ( + torch.tensor( + scale_rule.max_allowed_e2m1_value() * scale_rule.max_allowed_e4m3_value(), + dtype=x_amax.dtype, + device=x_amax.device, + ) + / x_amax + ) + x_block_scaled = torch.where( + x_scales.unsqueeze(1) != 0, + x_scale_blocks * (1 / (decode_scale * x_scales.to(x_amax.dtype).unsqueeze(1))), + 0, + ) + + return x_block_scaled, x_scales + + +def select_fouroversix( + x_scale_blocks: torch.Tensor, + x_block_scaled_6: torch.Tensor, + scales_6: torch.Tensor, + x_block_scaled_4: torch.Tensor, + scales_4: torch.Tensor, + x_amax: torch.Tensor, + *, + scale_rule: ScaleRule = ScaleRule.mse, + round_style: RoundStyle = RoundStyle.nearest, +) -> tuple[torch.Tensor, torch.Tensor]: + x_fake_quantized_6 = fake_quantize_to_e2m1( + x_block_scaled_6, + round_style=round_style, + ) + x_fake_quantized_4 = fake_quantize_to_e2m1( + x_block_scaled_4, + round_style=round_style, + ) + + x_dequantized_6 = ( + x_fake_quantized_6.to(x_amax.dtype) + * scales_6.unsqueeze(1).to(x_amax.dtype) + * x_amax + / torch.tensor( + E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX, + dtype=x_amax.dtype, + device=x_amax.device, + ) + ) + x_dequantized_4 = ( + x_fake_quantized_4.to(x_amax.dtype) + * scales_4.unsqueeze(1).to(x_amax.dtype) + * x_amax + / torch.tensor( + E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX, + dtype=x_amax.dtype, + device=x_amax.device, + ) + ) + + if scale_rule == ScaleRule.abs_max: + x_error_4 = (x_dequantized_4 - x_scale_blocks).abs().max(axis=-1).values + x_error_6 = (x_dequantized_6 - x_scale_blocks).abs().max(axis=-1).values + elif scale_rule == ScaleRule.mae: + x_error_4 = (x_dequantized_4 - x_scale_blocks).abs().sum(axis=-1) + x_error_6 = (x_dequantized_6 - x_scale_blocks).abs().sum(axis=-1) + elif scale_rule == ScaleRule.mse: + x_error_4 = ((x_dequantized_4 - x_scale_blocks) ** 2).sum(axis=-1) + x_error_6 = ((x_dequantized_6 - x_scale_blocks) ** 2).sum(axis=-1) + + select_4 = (x_error_4 < x_error_6).unsqueeze(1) + x_fake_quantized = torch.where( + select_4, + x_fake_quantized_4.reshape(x_scale_blocks.shape[0], -1), + x_fake_quantized_6.reshape(x_scale_blocks.shape[0], -1), + ) + scales = torch.where( + select_4, + scales_4.reshape(-1, 1), + scales_6.reshape(-1, 1), + ) + + return x_fake_quantized, scales + + +def quantize_to_fp4( + x: torch.Tensor, + x_amax: torch.Tensor | None = None, + had: torch.Tensor | None = None, + *, + block_scale_2d: bool = False, + fp4_format: DataType = DataType.nvfp4, + round_style: RoundStyle = RoundStyle.nearest, + scale_rule: ScaleRule = ScaleRule.mse, + transpose: bool = False, +) -> ( + tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] + | tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor] +): + if transpose: + x = x.T + + if had is not None: + x = (x.reshape(-1, had.shape[0]) @ had).reshape_as(x) + + if x_amax is None: + x_amax = ( + torch.ones(1, device=x.device, dtype=x.dtype) + if fp4_format == DataType.mxfp4 + else x.abs().max().float() + ) + + if block_scale_2d: + assert x.ndim == 2 # noqa: PLR2004 + assert x.shape[1] % fp4_format.block_size() == 0 + + x_scale_blocks = ( + x.reshape( + -1, + fp4_format.block_size(), + x.shape[1] // fp4_format.block_size(), + fp4_format.block_size(), + ) + .permute(0, 2, 1, 3) + .reshape(-1, fp4_format.block_size() ** 2) + .float() + ) + else: + x_scale_blocks = x.reshape(-1, fp4_format.block_size()).float() + + x_fake_quantized = None + + if fp4_format == DataType.mxfp4: + x_block_scaled, scales = quantize_to_mxfp4( + x_scale_blocks, + scale_rule=scale_rule, + ) + elif fp4_format == DataType.nvfp4 and scale_rule in { + ScaleRule.static_6, + ScaleRule.static_4, + }: + x_block_scaled, scales = quantize_to_nvfp4( + x_scale_blocks, + x_amax, + scale_rule=scale_rule, + ) + elif fp4_format == DataType.nvfp4: # Four over six + x_block_scaled_6, scales_6 = quantize_to_nvfp4( + x_scale_blocks, + x_amax, + scale_rule=scale_rule, + ) + x_block_scaled_4, scales_4 = quantize_to_nvfp4( + x_scale_blocks, + x_amax, + scale_rule=scale_rule, + scale_expansion_factor=1.5, + ) + x_fake_quantized, scales = select_fouroversix( + x_scale_blocks, + x_block_scaled_6, + scales_6, + x_block_scaled_4, + scales_4, + x_amax, + scale_rule=scale_rule, + round_style=round_style, + ) + else: + msg = f"Invalid FP4 format: {fp4_format}" + raise ValueError(msg) + + if x_fake_quantized is None: + x_fake_quantized = fake_quantize_to_e2m1( + x_block_scaled, + round_style=round_style, + ) + + if block_scale_2d: + x_fake_quantized = x_fake_quantized.reshape( + -1, + x.shape[1] // fp4_format.block_size(), + fp4_format.block_size(), + fp4_format.block_size(), + ).permute(0, 2, 1, 3) + + scales = ( + scales.reshape( + 1, + x.shape[0] // fp4_format.block_size(), + x.shape[1] // fp4_format.block_size(), + ) + .broadcast_to( + fp4_format.block_size(), + x.shape[0] // fp4_format.block_size(), + x.shape[1] // fp4_format.block_size(), + ) + .permute(1, 0, 2) + ) + + x_quantized = pack_unpacked_fp4( + quantize_bf16_to_unpacked_fp4(x_fake_quantized.bfloat16().reshape_as(x)), + ) + + reshaped_scales = to_blocked( + scales.reshape( + x.shape[0], + x.shape[1] // fp4_format.block_size(), + ), + ) + + return x_quantized, reshaped_scales, x_amax diff --git a/fouroversix/src/fouroversix/quantize/quantized_tensor.py b/fouroversix/src/fouroversix/quantize/quantized_tensor.py new file mode 100644 index 0000000..03b9a60 --- /dev/null +++ b/fouroversix/src/fouroversix/quantize/quantized_tensor.py @@ -0,0 +1,248 @@ +from dataclasses import dataclass + +import torch +import torch.nn.functional as F +from fouroversix.utils import DataType, ScaleRule + +from .utils import to_blocked + + +def from_blocked(a: torch.Tensor, orig_shape: tuple[int, int]) -> torch.Tensor: + rows, cols = orig_shape + return ( + a.view(-1, 32, 4, 4) + .transpose(1, 2) + .reshape(-1, cols // 4, 128, 4) + .transpose(1, 2) + .reshape(rows, cols) + ) + + +def convert_e2m1_to_fp8_e4m3(x: torch.Tensor) -> torch.Tensor: + sign = (x >> 3) & 0x1 + exponent = (x >> 1) & 0x3 + mantissa = x & 0x1 + + # Make adjustments + new_exponent = torch.where( + (exponent == 0) & (mantissa == 0), + 0, + (exponent + 6) & 0xF, + ) + new_mantissa = torch.where(exponent == 0, 0, mantissa << 2) + + return ((sign << 7) | (new_exponent << 3) | new_mantissa).view(torch.float8_e4m3fn) + + +def unpack_packed_fp4( + x: torch.Tensor, + to_dtype: torch.dtype = torch.float8_e4m3fn, +) -> torch.Tensor: + if to_dtype == torch.float8_e4m3fn: + convert_function = convert_e2m1_to_fp8_e4m3 + else: + msg = f"Unsupported dtype: {to_dtype}" + raise ValueError(msg) + + high = (x >> 4) & 0xF + low = x & 0xF + + return torch.stack( + [convert_function(low), convert_function(high)], + dim=-1, + ).reshape(x.shape[0], x.shape[1] * 2) + + +@dataclass +class QuantizedTensor: + """A quantized tensor.""" + + values: torch.Tensor + scale_factors: torch.Tensor + amax: torch.Tensor + + dtype: DataType + original_shape: tuple[int, int] + scale_rule: ScaleRule + + padded_shape: tuple[int, int] + + def __init__( + self, + values: torch.Tensor, + scale_factors: torch.Tensor, + amax: torch.Tensor, + dtype: DataType, + original_shape: tuple[int, int], + scale_rule: ScaleRule, + padded_shape: tuple[int, int] | None = None, + ) -> None: + super().__init__() + + if isinstance(dtype, str): + dtype = DataType(dtype) + + if isinstance(original_shape, torch.Size): + original_shape = tuple(original_shape) + + if isinstance(scale_rule, str): + scale_rule = ScaleRule(scale_rule) + + if isinstance(padded_shape, torch.Size): + padded_shape = tuple(padded_shape) + + self.dtype = dtype + self.original_shape = original_shape + self.scale_rule = scale_rule + self.padded_shape = padded_shape + + if self.padded_shape is None: + rows_div = 128 + + # The scale factor layout requires 4 blocks along the K dimension for both + # MXFP4 and NVFP4. See: + # https://docs.nvidia.com/cutlass/latest/media/docs/cpp/blackwell_functionality.html#scale-factor-layouts + cols_div = 4 * dtype.block_size() + + self.padded_shape = ( + original_shape[0] + + (rows_div - original_shape[0] % rows_div) % rows_div, + original_shape[1] + + (cols_div - original_shape[1] % cols_div) % cols_div, + ) + + expected_packed_elements = self.padded_shape[0] * self.padded_shape[1] // 2 + expected_scale_factors = expected_packed_elements * 2 // dtype.block_size() + + if values.numel() != expected_packed_elements: + values = F.pad( + values, + ( + 0, + # Divide by 2 because these are packed values + self.padded_shape[1] // 2 - values.shape[1], + 0, + self.padded_shape[0] - values.shape[0], + ), + ) + + # If the scale factors are 1D, we assume that they are already in the + # correct layout for Blackwell. See: + # https://docs.nvidia.com/cutlass/latest/media/docs/cpp/blackwell_functionality.html#scale-factor-layouts + if ( + scale_factors.ndim > 1 + and scale_factors.numel() != expected_scale_factors + ): + scale_factors = F.pad( + scale_factors, + ( + 0, + ( + self.padded_shape[1] // dtype.block_size() + - scale_factors.shape[1] + ), + 0, + self.padded_shape[0] - scale_factors.shape[0], + ), + value=0 if dtype == DataType.nvfp4 else 1, + ) + + scale_factors = to_blocked(scale_factors) + + if values.numel() != expected_packed_elements: + msg = ( + f"Expected {expected_packed_elements} e2m1 values, got " + f"{values.numel()}" + ) + raise ValueError(msg) + + if scale_factors.numel() != expected_scale_factors: + msg = ( + f"Expected {expected_scale_factors} scale factors, got " + f"{scale_factors.numel()}" + ) + raise ValueError(msg) + # preprocess scale_factors + # padded_shape = self.padded_shape + # scales_2d = from_blocked( + # scale_factors, + # (padded_shape[0], padded_shape[1] // 16), + # ) + + # # 2. 正确计算 global scale: amax / (max_e2m1 * max_e4m3) + # global_scale = amax / (6.0 * 256.0) + # global_scale = amax / ( + # self.scale_rule.max_allowed_e2m1_value() + # * self.scale_rule.max_allowed_e4m3_value() + # ) + + + self.values = values + self.scale_factors = scale_factors + self.amax = amax + + def dequantize(self, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor: + """Return a high-precision tensor with the dequantized values (PyTorch impl).""" + + values = unpack_packed_fp4(self.values).to(dtype) + scales = from_blocked( + self.scale_factors, + ( + self.padded_shape[0], + self.padded_shape[1] // self.dtype.block_size(), + ), + ) + + result = values * scales.to(dtype).repeat_interleave( + self.dtype.block_size(), + -1, + ) + + if self.dtype == DataType.nvfp4 and self.amax is not None: + result = ( + result.to(torch.float32) + * self.amax + / ( + self.scale_rule.max_allowed_e2m1_value() + * self.scale_rule.max_allowed_e4m3_value() + ) + ).to(dtype) + + if result.shape != self.original_shape: + result = result[: self.original_shape[0], : self.original_shape[1]] + + return result + + def dequantize_triton(self, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor: + """Return a high-precision tensor with the dequantized values (Triton kernel).""" + from utils.nvfp4_kernel import fp4_dequantize + + block_size = self.dtype.block_size() + + scales_2d = from_blocked( + self.scale_factors, + (self.padded_shape[0], self.padded_shape[1] // block_size), + ) + + global_scale = self.amax / ( + self.scale_rule.max_allowed_e2m1_value() + * self.scale_rule.max_allowed_e4m3_value() + ) + + result = fp4_dequantize( + self.values, + scales_2d, + global_scale, + block_size=block_size, + dtype=dtype, + ) + + if result.shape != self.original_shape: + result = result[: self.original_shape[0], : self.original_shape[1]] + + return result + + @property + def device(self) -> torch.device: + """Get device of the values in this tensor.""" + return self.values.device diff --git a/fouroversix/src/fouroversix/quantize/transformer_engine.py b/fouroversix/src/fouroversix/quantize/transformer_engine.py new file mode 100644 index 0000000..cd69b3d --- /dev/null +++ b/fouroversix/src/fouroversix/quantize/transformer_engine.py @@ -0,0 +1,109 @@ +import functools + +import torch +from fouroversix.quantize.utils import to_blocked +from fouroversix.utils import BLACKWELL_SM_IDS, DataType, RoundStyle, ScaleRule + +from .backend import QuantizeBackendBase +from .config import QuantizationConfig +from .quantized_tensor import QuantizedTensor + + +class TransformerEngineQuantizeBackend(QuantizeBackendBase): + """ + A backend that quantizes inputs using NVIDIA's TransformerEngine. Used to debug + our implementations. + """ + + @classmethod + @functools.lru_cache + def is_available(cls) -> bool: + """ + Return True if the Transformer Engine backend is available on the current + machine. + """ + + if ( + not torch.cuda.is_available() + or torch.cuda.get_device_capability()[0] not in BLACKWELL_SM_IDS + ): + return False + + try: + import transformer_engine # noqa: F401 + except ModuleNotFoundError: + return False + + return True + + @classmethod + def is_supported(cls, x: torch.Tensor, config: QuantizationConfig) -> bool: + """ + Return True if the Transformer Engine backend supports the given input and + quantization configuration. + """ + + if not super().is_supported(x, config): + return False + + if config.dtype != DataType.nvfp4 or config.scale_rule != ScaleRule.static_6: + return False + + if not config.transpose and config.rht: + return False + + if config.transpose and config.rht and config.block_scale_2d: # noqa: SIM103 + return False + + return True + + @classmethod + def quantize_to_fp4( + cls, + x: torch.Tensor, + config: QuantizationConfig, + ) -> QuantizedTensor: + """ + Quantize a tensor to FP4 using the Transformer Engine backend. + + Args: + x (torch.Tensor): The input tensor to quantize. + config (QuantizationConfig): The quantization configuration. + + Returns: + The quantized tensor. + + """ + + from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer + + q = NVFP4Quantizer( + with_2d_quantization=config.block_scale_2d, + with_rht=config.rht, + with_post_rht_amax=config.rht, + stochastic_rounding=config.round_style == RoundStyle.stochastic, + ) + + out = q.quantize(x) + + if config.transpose: + values = out._columnwise_data # noqa: SLF001 + scale_factors = to_blocked( + out._columnwise_scale_inv.view(torch.float8_e4m3fn), # noqa: SLF001 + ) + amax = out._amax_columnwise # noqa: SLF001 + else: + values = out._rowwise_data # noqa: SLF001 + scale_factors = to_blocked( + out._rowwise_scale_inv.view(torch.float8_e4m3fn), # noqa: SLF001 + ) + amax = out._amax_rowwise # noqa: SLF001 + + return QuantizedTensor( + values, + scale_factors, + amax, + config.dtype, + (x.shape[1], x.shape[0]) if config.transpose else x.shape, + config.scale_rule, + ) diff --git a/fouroversix/src/fouroversix/quantize/triton/__init__.py b/fouroversix/src/fouroversix/quantize/triton/__init__.py new file mode 100644 index 0000000..438e9d0 --- /dev/null +++ b/fouroversix/src/fouroversix/quantize/triton/__init__.py @@ -0,0 +1,3 @@ +from .backend import TritonQuantizeBackend + +__all__ = ["TritonQuantizeBackend"] diff --git a/fouroversix/src/fouroversix/quantize/triton/backend.py b/fouroversix/src/fouroversix/quantize/triton/backend.py new file mode 100644 index 0000000..a76c7df --- /dev/null +++ b/fouroversix/src/fouroversix/quantize/triton/backend.py @@ -0,0 +1,79 @@ +import functools + +import torch +from fouroversix.quantize.backend import QuantizeBackendBase +from fouroversix.quantize.config import QuantizationConfig +from fouroversix.quantize.quantized_tensor import QuantizedTensor +from fouroversix.quantize.utils import get_rht_matrix +from fouroversix.utils import BLACKWELL_SM_IDS, SM_100, RoundStyle + + +class TritonQuantizeBackend(QuantizeBackendBase): + """ + The Triton quantization backend. Supports all parameters required for efficient + NVFP4 training, including stochastic rounding, the random Hadamard transform, + transposed inputs, and 2D block scaling. Requires a Blackwell GPU. + """ + + @classmethod + @functools.lru_cache + def is_available(cls) -> bool: + """Return True if the Triton backend is available on the current machine.""" + return ( + torch.cuda.is_available() + and torch.cuda.get_device_capability()[0] in BLACKWELL_SM_IDS + ) + + @classmethod + def is_supported(cls, x: torch.Tensor, config: QuantizationConfig) -> bool: + """ + Return True if the Triton backend supports the given input and quantization + configuration. + """ + + if not super().is_supported(x, config): + return False + + if config.round_style == RoundStyle.stochastic: + return torch.cuda.get_device_capability()[0] == SM_100 + + return x.device.type == "cuda" + + @classmethod + def quantize_to_fp4( + cls, + x: torch.Tensor, + config: QuantizationConfig, + ) -> QuantizedTensor: + """ + Quantize a tensor to FP4 using the Triton backend. + + Args: + x (torch.Tensor): The input tensor to quantize. + config (QuantizationConfig): The quantization configuration. + + Returns: + The quantized tensor. + + """ + + from .kernel import quantize_to_fp4 + + values, scale_factors, amax = quantize_to_fp4( + x, + had=get_rht_matrix() if config.rht else None, + fp4_format=config.dtype, + round_style=config.round_style, + scale_rule=config.scale_rule, + block_scale_2d=config.block_scale_2d, + transpose=config.transpose, + ) + + return QuantizedTensor( + values, + scale_factors, + amax, + config.dtype, + (x.shape[1], x.shape[0]) if config.transpose else x.shape, + config.scale_rule, + ) diff --git a/fouroversix/src/fouroversix/quantize/triton/kernel.py b/fouroversix/src/fouroversix/quantize/triton/kernel.py new file mode 100644 index 0000000..f47efb9 --- /dev/null +++ b/fouroversix/src/fouroversix/quantize/triton/kernel.py @@ -0,0 +1,672 @@ +from __future__ import annotations + +import torch +import triton +import triton.language as tl +from fouroversix.utils import DataType, RoundStyle, ScaleRule +from triton.tools.tensor_descriptor import TensorDescriptor + +E2M1_MAX_VALUE = tl.constexpr(6) +E2M1_MAX_FOUR = tl.constexpr(4) +E4M3_MAX_VALUE = tl.constexpr(448) +E4M3_MAX_FOUROVERSIX = tl.constexpr(256) +SCALE_MEGABLOCK_SIZE = tl.constexpr(512) + +DATA_TYPE_MXFP4 = tl.constexpr(DataType.mxfp4.value) +DATA_TYPE_NVFP4 = tl.constexpr(DataType.nvfp4.value) + +ROUND_STYLE_NEAREST = tl.constexpr(RoundStyle.nearest.value) +ROUND_STYLE_STOCHASTIC = tl.constexpr(RoundStyle.stochastic.value) + +SCALE_RULE_ABS_MAX = tl.constexpr(ScaleRule.abs_max.value) +SCALE_RULE_STATIC_4 = tl.constexpr(ScaleRule.static_4.value) +SCALE_RULE_STATIC_6 = tl.constexpr(ScaleRule.static_6.value) +SCALE_RULE_MAE = tl.constexpr(ScaleRule.mae.value) +SCALE_RULE_MSE = tl.constexpr(ScaleRule.mse.value) + + +@triton.jit +def rht_kernel( + x_desc, + h_desc, + y_desc, + # Meta-parameters + # TODO(jack): Update RHT kernel to support unpadded dimensions + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + TRANSPOSE: tl.constexpr, +) -> None: + HAD_BLOCK_SIZE: tl.constexpr = h_desc.block_shape[0] + + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + # Load H [B, B] + h_block = h_desc.load([0, 0]) + + m_block_offset = pid_m * BLOCK_SIZE_M + n_block_offset = pid_n * BLOCK_SIZE_N + + if not TRANSPOSE: + x_block = x_desc.load([m_block_offset, n_block_offset]) + else: + x_block = x_desc.load([n_block_offset, m_block_offset]).T + + y_block = tl.dot( + x_block.reshape( + BLOCK_SIZE_M * BLOCK_SIZE_N // HAD_BLOCK_SIZE, + HAD_BLOCK_SIZE, + ).to(tl.bfloat16), + h_block, + ).reshape(BLOCK_SIZE_M, BLOCK_SIZE_N) + + y_desc.store([m_block_offset, n_block_offset], y_block) + + +@triton.jit +def block_scaled_fp4_quantization_kernel( + x_block, + x_amax_ptr, + rbits_ptr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + FP4_FORMAT: tl.constexpr, + ROUND_STYLE: tl.constexpr, + BLOCK_SCALE_2D: tl.constexpr, + SCALE_RULE: tl.constexpr, +) -> None: + E2M1_MAX_ALLOWED_VALUE: tl.constexpr = ( + E2M1_MAX_VALUE if SCALE_RULE == SCALE_RULE_STATIC_6 else E2M1_MAX_FOUR + ) + + if FP4_FORMAT == DATA_TYPE_MXFP4: + x_scale_blocks = x_block.reshape(128, 4, 32) + x_scales_hp = tl.max(x_scale_blocks.abs(), axis=-1) / E2M1_MAX_ALLOWED_VALUE + x_scales_e8m0_u32 = x_scales_hp.cast(tl.uint32, bitcast=True) + + # Use the 8-bit exponent as the scale factor + x_scales_e8m0 = ((x_scales_e8m0_u32 >> 23) & 0xFF).to(tl.uint8) + + # Add one in order to round up + x_scales = tl.where( + (x_scales_e8m0_u32 & 0x7FFFFF) == 0, + x_scales_e8m0, + x_scales_e8m0 + 1, + ) + + # Convert the rounded-up scale factor back to a 32-bit float + x_scales_hp = (x_scales.cast(tl.uint32) << 23).cast(x_block.dtype, bitcast=True) + + if BLOCK_SCALE_2D: + x_scales_hp = ( + tl.max( + x_scales_hp.reshape(4, 32, 4).permute(0, 2, 1), + axis=-1, + ) + .expand_dims(0) + .broadcast_to(4, 32, 4) + .permute(1, 0, 2) + .reshape(128, 4) + ) + + (x_block_scaled_b1, x_block_scaled_b2) = ( + (x_scale_blocks / x_scales_hp.expand_dims(2)) + .reshape(BLOCK_SIZE_M, BLOCK_SIZE_N // 2, 2) + .split() + ) + elif FP4_FORMAT == DATA_TYPE_NVFP4: + x_amax = tl.load(x_amax_ptr) + x_scale_blocks = x_block.reshape(128, 4, 16) + + if x_amax == 0: + x_scales_hp = tl.full((128, 4), 0, dtype=tl.float32) + else: + encode_scale = tl.div_rn(E2M1_MAX_ALLOWED_VALUE * E4M3_MAX_VALUE, x_amax) + x_scales_hp = ( + tl.div_rn(tl.max(x_scale_blocks.abs(), axis=-1), E2M1_MAX_ALLOWED_VALUE) + * encode_scale + ) + + if BLOCK_SCALE_2D: + x_scales_hp = ( + tl.max( + x_scales_hp.reshape(8, 16, 4).permute(0, 2, 1), + axis=-1, + ) + .expand_dims(0) + .broadcast_to(16, 8, 4) + .permute(1, 0, 2) + .reshape(128, 4) + ) + + x_scales = x_scales_hp.to(tl.float8e4nv) + + decode_scale = tl.div_rn( + 1, + tl.div_rn(E2M1_MAX_ALLOWED_VALUE * E4M3_MAX_VALUE, x_amax), + ) + (x_block_scaled_b1, x_block_scaled_b2) = ( + tl.where( + x_scales.expand_dims(2).to(x_amax.dtype) != 0, + x_scale_blocks + * tl.div_rn(1, decode_scale * x_scales.to(x_amax.dtype).expand_dims(2)), + 0, + ) + .reshape(BLOCK_SIZE_M, BLOCK_SIZE_N // 2, 2) + .split() + ) + + if ROUND_STYLE == ROUND_STYLE_NEAREST: + x_e2m1 = tl.inline_asm_elementwise( + asm=""" + { + .reg .b8 byte0, byte1, byte2, byte3; + cvt.rn.satfinite.e2m1x2.f32 byte0, $5, $1; + cvt.rn.satfinite.e2m1x2.f32 byte1, $6, $2; + cvt.rn.satfinite.e2m1x2.f32 byte2, $7, $3; + cvt.rn.satfinite.e2m1x2.f32 byte3, $8, $4; + mov.b32 $0, {byte0, byte1, byte2, byte3}; + } + """, + constraints="=r,r,r,r,r,r,r,r,r", + args=[x_block_scaled_b1, x_block_scaled_b2], + dtype=tl.uint8, + is_pure=True, + pack=4, + ) + elif ROUND_STYLE == ROUND_STYLE_STOCHASTIC: + rbits = tl.randint( + tl.load(rbits_ptr), + tl.arange(0, BLOCK_SIZE_M)[:, None] * BLOCK_SIZE_N // 2 + + tl.arange(0, BLOCK_SIZE_N // 2)[None, :], + ).cast(tl.uint32, bitcast=True) + + x_e2m1 = tl.inline_asm_elementwise( + asm=""" + { + .reg .b16 tmp0, tmp1; + cvt.rs.satfinite.e2m1x4.f32 tmp0, {$6, $2, $5, $1}, $9; + cvt.rs.satfinite.e2m1x4.f32 tmp1, {$8, $4, $7, $3}, $10; + mov.b32 $0, {tmp0, tmp1}; + } + """, + constraints="=r,r,r,r,r,r,r,r,r,r,r,r,r", + args=[x_block_scaled_b1, x_block_scaled_b2, rbits], + dtype=tl.uint8, + is_pure=True, + pack=4, + ) + + return x_e2m1, x_scales.reshape(4, 32, 4).permute(1, 0, 2).ravel() + + +@triton.jit +def nvfp4_fouroversix_quantization_kernel( # noqa: C901, PLR0915 + x_block, + x_amax_ptr, + rbits_ptr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + ROUND_STYLE: tl.constexpr, + BLOCK_SCALE_2D: tl.constexpr, + SCALE_RULE: tl.constexpr, +) -> None: + x_amax = tl.load(x_amax_ptr) + x_scale_blocks = x_block.reshape(128, 4, 16) + + if x_amax == 0: + x_scales_hp = tl.full((128, 4), 0, dtype=tl.float32) + else: + encode_scale = tl.div_rn(E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX, x_amax) + x_scales_hp = ( + tl.div_rn(tl.max(x_scale_blocks.abs(), axis=-1), E2M1_MAX_VALUE) + * encode_scale + ) + + if BLOCK_SCALE_2D: + x_scales_hp = ( + tl.max( + x_scales_hp.reshape(8, 16, 4).permute(0, 2, 1), + axis=-1, + ) + .expand_dims(0) + .broadcast_to(16, 8, 4) + .permute(1, 0, 2) + .reshape(128, 4) + ) + + x_scales_6 = x_scales_hp.to(tl.float8e4nv) + x_scales_4 = (x_scales_hp * 1.5).to(tl.float8e4nv) + + decode_scale = tl.div_rn( + 1, + tl.div_rn(E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX, x_amax), + ) + + (x_block_scaled_6_b1, x_block_scaled_6_b2) = ( + tl.where( + x_scales_6.expand_dims(2).to(x_amax.dtype) != 0, + x_scale_blocks + * tl.div_rn(1, decode_scale * x_scales_6.to(x_amax.dtype).expand_dims(2)), + 0, + ) + .reshape(BLOCK_SIZE_M, BLOCK_SIZE_N // 2, 2) + .split() + ) + + (x_block_scaled_4_b1, x_block_scaled_4_b2) = ( + tl.where( + x_scales_4.expand_dims(2).to(x_amax.dtype) != 0, + x_scale_blocks + * tl.div_rn(1, decode_scale * x_scales_4.to(x_amax.dtype).expand_dims(2)), + 0, + ) + .reshape(BLOCK_SIZE_M, BLOCK_SIZE_N // 2, 2) + .split() + ) + + if ROUND_STYLE == ROUND_STYLE_NEAREST: + (x_e2m1_6, x_e2m1_4, x_fp16x2_6, x_fp16x2_4) = tl.inline_asm_elementwise( + asm=""" + { + .reg .b8 byte0, byte1, byte2, byte3; + + cvt.rn.satfinite.e2m1x2.f32 byte0, $28, $20; + cvt.rn.f16x2.e2m1x2 $4, byte0; + cvt.rn.satfinite.e2m1x2.f32 byte1, $29, $21; + cvt.rn.f16x2.e2m1x2 $5, byte1; + cvt.rn.satfinite.e2m1x2.f32 byte2, $30, $22; + cvt.rn.f16x2.e2m1x2 $6, byte2; + cvt.rn.satfinite.e2m1x2.f32 byte3, $31, $23; + cvt.rn.f16x2.e2m1x2 $7, byte3; + mov.b32 $0, {byte0, byte1, byte2, byte3}; + + cvt.rn.satfinite.e2m1x2.f32 byte0, $32, $24; + cvt.rn.f16x2.e2m1x2 $8, byte0; + cvt.rn.satfinite.e2m1x2.f32 byte1, $33, $25; + cvt.rn.f16x2.e2m1x2 $9, byte1; + cvt.rn.satfinite.e2m1x2.f32 byte2, $34, $26; + cvt.rn.f16x2.e2m1x2 $10, byte2; + cvt.rn.satfinite.e2m1x2.f32 byte3, $35, $27; + cvt.rn.f16x2.e2m1x2 $11, byte3; + mov.b32 $1, {byte0, byte1, byte2, byte3}; + + cvt.rn.satfinite.e2m1x2.f32 byte0, $44, $36; + cvt.rn.f16x2.e2m1x2 $12, byte0; + cvt.rn.satfinite.e2m1x2.f32 byte1, $45, $37; + cvt.rn.f16x2.e2m1x2 $13, byte1; + cvt.rn.satfinite.e2m1x2.f32 byte2, $46, $38; + cvt.rn.f16x2.e2m1x2 $14, byte2; + cvt.rn.satfinite.e2m1x2.f32 byte3, $47, $39; + cvt.rn.f16x2.e2m1x2 $15, byte3; + mov.b32 $2, {byte0, byte1, byte2, byte3}; + + cvt.rn.satfinite.e2m1x2.f32 byte0, $48, $40; + cvt.rn.f16x2.e2m1x2 $16, byte0; + cvt.rn.satfinite.e2m1x2.f32 byte1, $49, $41; + cvt.rn.f16x2.e2m1x2 $17, byte1; + cvt.rn.satfinite.e2m1x2.f32 byte2, $50, $42; + cvt.rn.f16x2.e2m1x2 $18, byte2; + cvt.rn.satfinite.e2m1x2.f32 byte3, $51, $43; + cvt.rn.f16x2.e2m1x2 $19, byte3; + mov.b32 $3, {byte0, byte1, byte2, byte3}; + } + """, + constraints="=r,=r,=r,=r,=r,=r,=r,=r,=r,=r,=r,=r,=r,=r,=r,=r,=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r", + args=[ + x_block_scaled_6_b1, + x_block_scaled_6_b2, + x_block_scaled_4_b1, + x_block_scaled_4_b2, + ], + dtype=(tl.uint8, tl.uint8, tl.uint32, tl.uint32), + is_pure=True, + pack=8, + ) + elif ROUND_STYLE == ROUND_STYLE_STOCHASTIC: + rbits = tl.randint( + tl.load(rbits_ptr), + tl.arange(0, BLOCK_SIZE_M)[:, None] * BLOCK_SIZE_N // 2 + + tl.arange(0, BLOCK_SIZE_N // 2)[None, :], + ).cast(tl.uint32, bitcast=True) + + (x_e2m1_6, x_e2m1_4, x_fp16x2_6, x_fp16x2_4) = tl.inline_asm_elementwise( + asm=""" + { + .reg .b16 tmp0, tmp1; + .reg .b8 byte0, byte1; + + cvt.rs.satfinite.e2m1x4.f32 tmp0, {$29, $21, $28, $20}, $52; + mov.b16 {byte1, byte0}, tmp0; + cvt.rn.f16x2.e2m1x2 $4, byte0; + cvt.rn.f16x2.e2m1x2 $5, byte1; + cvt.rs.satfinite.e2m1x4.f32 tmp1, {$31, $23, $30, $22}, $53; + mov.b16 {byte1, byte0}, tmp1; + cvt.rn.f16x2.e2m1x2 $6, byte0; + cvt.rn.f16x2.e2m1x2 $7, byte1; + mov.b32 $0, {tmp0, tmp1}; + + cvt.rs.satfinite.e2m1x4.f32 tmp0, {$33, $25, $32, $24}, $54; + mov.b16 {byte1, byte0}, tmp0; + cvt.rn.f16x2.e2m1x2 $8, byte0; + cvt.rn.f16x2.e2m1x2 $9, byte1; + cvt.rs.satfinite.e2m1x4.f32 tmp1, {$35, $27, $34, $26}, $55; + mov.b16 {byte1, byte0}, tmp1; + cvt.rn.f16x2.e2m1x2 $10, byte0; + cvt.rn.f16x2.e2m1x2 $11, byte1; + mov.b32 $1, {tmp0, tmp1}; + + cvt.rs.satfinite.e2m1x4.f32 tmp0, {$45, $37, $44, $36}, $56; + mov.b16 {byte1, byte0}, tmp0; + cvt.rn.f16x2.e2m1x2 $12, byte0; + cvt.rn.f16x2.e2m1x2 $13, byte1; + cvt.rs.satfinite.e2m1x4.f32 tmp1, {$47, $39, $46, $38}, $57; + mov.b16 {byte1, byte0}, tmp1; + cvt.rn.f16x2.e2m1x2 $14, byte0; + cvt.rn.f16x2.e2m1x2 $15, byte1; + mov.b32 $2, {tmp0, tmp1}; + + cvt.rs.satfinite.e2m1x4.f32 tmp0, {$49, $41, $48, $40}, $58; + mov.b16 {byte1, byte0}, tmp0; + cvt.rn.f16x2.e2m1x2 $16, byte0; + cvt.rn.f16x2.e2m1x2 $17, byte1; + cvt.rs.satfinite.e2m1x4.f32 tmp1, {$51, $43, $50, $42}, $59; + mov.b16 {byte1, byte0}, tmp1; + cvt.rn.f16x2.e2m1x2 $18, byte0; + cvt.rn.f16x2.e2m1x2 $19, byte1; + mov.b32 $3, {tmp0, tmp1}; + } + """, + constraints="=r,=r,=r,=r,=r,=r,=r,=r,=r,=r,=r,=r,=r,=r,=r,=r,=r,=r,=r,=r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r", + args=[ + x_block_scaled_6_b1, + x_block_scaled_6_b2, + x_block_scaled_4_b1, + x_block_scaled_4_b2, + rbits, + ], + dtype=(tl.uint8, tl.uint8, tl.uint32, tl.uint32), + is_pure=True, + pack=8, + ) + + x_fp16_6_lo = ( + (x_fp16x2_6 & 0xFFFF) + .cast(tl.uint16) + .cast(tl.float16, bitcast=True) + .cast(x_amax.dtype) + ) + x_fp16_6_hi = ( + (x_fp16x2_6 >> 16) + .cast(tl.uint16) + .cast(tl.float16, bitcast=True) + .cast(x_amax.dtype) + ) + x_hp_6 = tl.join(x_fp16_6_lo, x_fp16_6_hi).reshape(128, 4, 16) + + x_dequantized_6 = tl.div_rn( + x_hp_6 * x_scales_6.to(x_amax.dtype).expand_dims(2) * x_amax, + E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX, + ) + + x_fp16_4_lo = ( + (x_fp16x2_4 & 0xFFFF) + .cast(tl.uint16) + .cast(tl.float16, bitcast=True) + .cast(x_amax.dtype) + ) + x_fp16_4_hi = ( + (x_fp16x2_4 >> 16) + .cast(tl.uint16) + .cast(tl.float16, bitcast=True) + .cast(x_amax.dtype) + ) + x_hp_4 = tl.join(x_fp16_4_lo, x_fp16_4_hi).reshape(128, 4, 16) + + x_dequantized_4 = tl.div_rn( + x_hp_4 * x_scales_4.to(x_amax.dtype).expand_dims(2) * x_amax, + E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX, + ) + + diff_6 = x_dequantized_6 - x_scale_blocks + diff_4 = x_dequantized_4 - x_scale_blocks + + if SCALE_RULE == SCALE_RULE_ABS_MAX: + six_error = tl.max(tl.abs(diff_6), axis=-1) + four_error = tl.max(tl.abs(diff_4), axis=-1) + elif SCALE_RULE == SCALE_RULE_MAE: + six_error = tl.sum(tl.abs(diff_6), axis=-1) + four_error = tl.sum(tl.abs(diff_4), axis=-1) + elif SCALE_RULE == SCALE_RULE_MSE: + six_error = tl.sum(diff_6 * diff_6, axis=-1) + four_error = tl.sum(diff_4 * diff_4, axis=-1) + + if BLOCK_SCALE_2D: + six_error = six_error.reshape(8, 16, 4).permute(0, 2, 1) + four_error = four_error.reshape(8, 16, 4).permute(0, 2, 1) + + if SCALE_RULE == SCALE_RULE_ABS_MAX: + six_error = tl.max(six_error, axis=-1) + four_error = tl.max(four_error, axis=-1) + elif SCALE_RULE == SCALE_RULE_MAE or SCALE_RULE == SCALE_RULE_MSE: + six_error = tl.sum(six_error, axis=-1) + four_error = tl.sum(four_error, axis=-1) + + six_error = ( + six_error.expand_dims(0) + .broadcast_to(16, 8, 4) + .permute(1, 0, 2) + .reshape(128, 4) + ) + four_error = ( + four_error.expand_dims(0) + .broadcast_to(16, 8, 4) + .permute(1, 0, 2) + .reshape(128, 4) + ) + + x_e2m1 = tl.where( + (four_error < six_error).expand_dims(2), + x_e2m1_4.reshape(128, 4, 8), + x_e2m1_6.reshape(128, 4, 8), + ).reshape(128, 32) + + x_scales = ( + tl.where(four_error < six_error, x_scales_4, x_scales_6) + .reshape(4, 32, 4) + .permute(1, 0, 2) + .ravel() + ) + + return x_e2m1, x_scales + + +@triton.jit +def fp4_quantization_kernel( + x_desc, + x_amax_ptr, + x_e2m1_desc, + x_sf_desc, + rbits_ptr, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + TRANSPOSE: tl.constexpr, + FP4_FORMAT: tl.constexpr, + ROUND_STYLE: tl.constexpr, + BLOCK_SCALE_2D: tl.constexpr, + SCALE_RULE: tl.constexpr, +) -> None: + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + m_block_offset = pid_m * BLOCK_SIZE_M + n_block_offset = pid_n * BLOCK_SIZE_N + + # Load [B, B] block from A or A^T + if not TRANSPOSE: + x_block = x_desc.load([m_block_offset, n_block_offset]) + else: + x_block = x_desc.load([n_block_offset, m_block_offset]).T + + x_block = x_block.to(tl.float32) + + if SCALE_RULE == SCALE_RULE_STATIC_6 or SCALE_RULE == SCALE_RULE_STATIC_4: + x_e2m1, x_scales = block_scaled_fp4_quantization_kernel( + x_block, + x_amax_ptr, + rbits_ptr, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + FP4_FORMAT, + ROUND_STYLE, + BLOCK_SCALE_2D, + SCALE_RULE, + ) + else: + x_e2m1, x_scales = nvfp4_fouroversix_quantization_kernel( + x_block, + x_amax_ptr, + rbits_ptr, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + ROUND_STYLE, + BLOCK_SCALE_2D, + SCALE_RULE, + ) + + e2m1_n_block_offset = pid_n * BLOCK_SIZE_N // 2 + x_e2m1_desc.store([m_block_offset, e2m1_n_block_offset], x_e2m1) + + scale_block_offset = (pid_m * tl.num_programs(1) + pid_n) * SCALE_MEGABLOCK_SIZE + x_sf_desc.store([scale_block_offset], x_scales) + + +def quantize_to_fp4( + x: torch.Tensor, + x_amax: torch.Tensor | None = None, + had: torch.Tensor | None = None, + *, + fp4_format: DataType = DataType.nvfp4, + round_style: RoundStyle = RoundStyle.nearest, + scale_rule: ScaleRule = ScaleRule.mse, + block_scale_2d: bool = False, + transpose: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + if transpose: + N, M = x.shape + else: + M, N = x.shape + + block_size_m = 128 + block_size_n = 4 * fp4_format.block_size() + scale_dtype = torch.float8_e4m3fn if fp4_format == DataType.nvfp4 else torch.uint8 + + if x_amax is None: + x_amax = ( + x.abs().max().float() + if fp4_format == DataType.nvfp4 + else torch.ones(1, device=x.device, dtype=torch.float32) + ) + + padded_m = M + (block_size_m - M % block_size_m) % block_size_m + padded_n = N + (block_size_n - N % block_size_n) % block_size_n + + x_e2m1 = torch.empty((padded_m, padded_n // 2), device=x.device, dtype=torch.uint8) + x_sf = torch.empty( + padded_m * padded_n // fp4_format.block_size(), + device=x.device, + dtype=scale_dtype, + ) + + grid = lambda _: ( # noqa: E731 + padded_m // block_size_m, + padded_n // block_size_n, + ) + + x_desc = TensorDescriptor.from_tensor( + x, + block_shape=[ + block_size_m if not transpose else block_size_n, + block_size_n if not transpose else block_size_m, + ], + ) + x_e2m1_desc = TensorDescriptor.from_tensor( + x_e2m1, + block_shape=[block_size_m, block_size_n // 2], + ) + x_sf_desc = TensorDescriptor.from_tensor( + x_sf, + block_shape=[SCALE_MEGABLOCK_SIZE.value], + ) + + if had is not None: + had_block_size = had.shape[0] + + if M % had_block_size != 0: + msg = ( + f"The first dimension of A ({M}) must be divisible by the width of H " + f"({had_block_size})" + ) + raise ValueError(msg) + if N % had_block_size != 0: + msg = ( + f"The second dimension of A ({N}) must be divisible by the width of H " + f"({had_block_size})" + ) + raise ValueError(msg) + if had.shape[0] != had.shape[1]: + msg = "H must be a square matrix" + raise ValueError(msg) + if (had.shape[0] & (had.shape[0] - 1)) != 0: + msg = "H must have dimensions that are a power of two" + raise ValueError(msg) + + x_rht = torch.empty((M, N), device=x.device, dtype=torch.bfloat16) + + h_desc = TensorDescriptor.from_tensor( + had, + block_shape=[had_block_size, had_block_size], + ) + x_rht_desc = TensorDescriptor.from_tensor( + x_rht, + block_shape=[block_size_m, block_size_n], + ) + + rht_kernel[grid]( + x_desc, + h_desc, + x_rht_desc, + BLOCK_SIZE_M=block_size_m, + BLOCK_SIZE_N=block_size_n, + TRANSPOSE=transpose, + ) + + transpose = False + x_amax = x_rht.abs().max().float() + + rbits = ( + torch.randint(0, torch.iinfo(torch.int32).max, (1,), device=x.device) + if round_style == RoundStyle.stochastic + else None + ) + + fp4_quantization_kernel[grid]( + x_rht_desc if had is not None else x_desc, + x_amax, + x_e2m1_desc, + x_sf_desc, + rbits, + BLOCK_SIZE_M=block_size_m, + BLOCK_SIZE_N=block_size_n, + TRANSPOSE=transpose, + FP4_FORMAT=fp4_format.value, + ROUND_STYLE=round_style.value, + BLOCK_SCALE_2D=block_scale_2d, + SCALE_RULE=scale_rule.value, + ) + + if fp4_format == DataType.mxfp4: + x_sf = x_sf.view(torch.float8_e8m0fnu) + + return x_e2m1, x_sf, x_amax diff --git a/fouroversix/src/fouroversix/quantize/utils.py b/fouroversix/src/fouroversix/quantize/utils.py new file mode 100644 index 0000000..a7720a4 --- /dev/null +++ b/fouroversix/src/fouroversix/quantize/utils.py @@ -0,0 +1,96 @@ +import functools +import math + +import torch + +""" +Credit: TransformerEngine +https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/tensor/nvfp4_tensor.py +""" + +HADAMARD_DIMENSION = 16 + + +def get_no_random_sign_vector(device: int) -> torch.Tensor: + """Non-random sign vector for Hadamard transform.""" + return torch.tensor([1], dtype=torch.float32, device=device) + + +def get_wgrad_sign_vector(device: int) -> torch.Tensor: + """ + Hard-coded random signs for Hadamard transform. + + https://xkcd.com/221/ + + """ + return torch.tensor( + [1, 1, 1, -1, 1, -1, -1, -1, -1, -1, -1, 1, -1, 1, -1, -1], + dtype=torch.float32, + device=device, + ) + + +def get_hadamard_matrix(hadamard_dimension: int, device: int) -> torch.Tensor: + """Construct a 16x16 Hadamard matrix.""" + + if hadamard_dimension != HADAMARD_DIMENSION: + msg = f"Only hadamard dimension {HADAMARD_DIMENSION} is supported." + raise ValueError(msg) + + hadamard_scale = 1 / math.sqrt(hadamard_dimension) + return ( + torch.tensor( + [ + [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], + [1, -1, 1, -1, 1, -1, 1, -1, 1, -1, 1, -1, 1, -1, 1, -1], + [1, 1, -1, -1, 1, 1, -1, -1, 1, 1, -1, -1, 1, 1, -1, -1], + [1, -1, -1, 1, 1, -1, -1, 1, 1, -1, -1, 1, 1, -1, -1, 1], + [1, 1, 1, 1, -1, -1, -1, -1, 1, 1, 1, 1, -1, -1, -1, -1], + [1, -1, 1, -1, -1, 1, -1, 1, 1, -1, 1, -1, -1, 1, -1, 1], + [1, 1, -1, -1, -1, -1, 1, 1, 1, 1, -1, -1, -1, -1, 1, 1], + [1, -1, -1, 1, -1, 1, 1, -1, 1, -1, -1, 1, -1, 1, 1, -1], + [1, 1, 1, 1, 1, 1, 1, 1, -1, -1, -1, -1, -1, -1, -1, -1], + [1, -1, 1, -1, 1, -1, 1, -1, -1, 1, -1, 1, -1, 1, -1, 1], + [1, 1, -1, -1, 1, 1, -1, -1, -1, -1, 1, 1, -1, -1, 1, 1], + [1, -1, -1, 1, 1, -1, -1, 1, -1, 1, 1, -1, -1, 1, 1, -1], + [1, 1, 1, 1, -1, -1, -1, -1, -1, -1, -1, -1, 1, 1, 1, 1], + [1, -1, 1, -1, -1, 1, -1, 1, -1, 1, -1, 1, 1, -1, 1, -1], + [1, 1, -1, -1, -1, -1, 1, 1, -1, -1, 1, 1, 1, 1, -1, -1], + [1, -1, -1, 1, -1, 1, 1, -1, -1, 1, 1, -1, 1, -1, -1, 1], + ], + dtype=torch.float32, + device=device, + ) + * hadamard_scale + ) + + +@functools.cache +def get_rht_matrix( + *, + with_random_sign_mask: bool = True, + device: str | int = "cuda", +) -> torch.Tensor: + """Construct matrix used in random Hadamard transform.""" + if with_random_sign_mask: + signs = get_wgrad_sign_vector(device=device) + else: + signs = get_no_random_sign_vector(device=device) + sign_matrix = signs * torch.eye( + HADAMARD_DIMENSION, + dtype=torch.float32, + device=device, + ) + rht_matrix = sign_matrix @ get_hadamard_matrix(HADAMARD_DIMENSION, device=device) + return rht_matrix.to(dtype=torch.bfloat16) + + +def to_blocked(a: torch.Tensor) -> torch.Tensor: + return ( + a.view(a.shape[0] // 128, 128, a.shape[1] // 4, 4) + .transpose(1, 2) + .reshape(-1, 4, 32, 4) + .transpose(1, 2) + .reshape(-1, 32, 16) + .flatten() + ) diff --git a/fouroversix/src/fouroversix/utils.py b/fouroversix/src/fouroversix/utils.py new file mode 100644 index 0000000..c1516ba --- /dev/null +++ b/fouroversix/src/fouroversix/utils.py @@ -0,0 +1,134 @@ +from enum import Enum + +import torch + +SM_100 = 10 +SM_110 = 11 +SM_120 = 12 + +BLACKWELL_SM_IDS = {SM_100, SM_110, SM_120} + + +class DataType(str, Enum): + """Data types.""" + + bfloat16 = "bfloat16" + float16 = "float16" + float32 = "float32" + mxfp4 = "mxfp4" + nvfp4 = "nvfp4" + + def block_size(self) -> int | None: + """Return the block size if this a block-scaled format, or `None` otherwise.""" + + return { + DataType.mxfp4: 32, + DataType.nvfp4: 16, + }.get(self) + + def scale_dtype(self) -> torch.dtype | None: + """Return the scale dtype if this a block-scaled format, or `None` otherwise.""" + + return { + DataType.mxfp4: torch.float8_e8m0fnu, + DataType.nvfp4: torch.float8_e4m3fn, + }.get(self) + + def torch_dtype(self) -> torch.dtype | None: + """ + Return the corresponding torch.dtype if one is available, or `None` + otherwise. + """ + + return { + DataType.bfloat16: torch.bfloat16, + DataType.float16: torch.float16, + DataType.float32: torch.float32, + }.get(self) + + +class MatmulBackend(str, Enum): + """ + Backends for matrix multiplication with FP4. + + - `cutlass`: CUTLASS implementation. This requires a Blackwell GPU. + - `pytorch`: PyTorch implementation which first dequantizes the input tensors to + FP32 and then performs an FP32 matrix multiplication. + """ + + cutlass = "cutlass" + pytorch = "pytorch" + + +class QuantizeBackend(str, Enum): + """ + Backends for quantizing a tensor to NVFP4 or MXFP4. + + - `cuda`: CUDA implementation. Requires a Blackwell GPU, and currently only supports + the forward pass for PTQ (no stochastic rounding, no transposed matrices, no + RHT, no 2D block scaling). + - `pytorch`: PyTorch implementation. + - `triton`: Triton implementation. Requires a Blackwell GPU. + """ + + cuda = "cuda" + pytorch = "pytorch" + transformer_engine = "transformer_engine" + triton = "triton" + + +class RoundStyle(str, Enum): + """ + Rounding styles for quantization. + + - `nearest`: Round to the nearest FP4 value. + - `stochastic`: Round to the nearest FP4 value after applying random noise to each + value. + """ + + nearest = "nearest" + stochastic = "stochastic" + + +class ScaleRule(str, Enum): + """ + Block scale selection rules for NVFP4 quantization. + + - `abs_max`: Between 4 and 6, select the block scale that minimizes the maximum + absolute quantization error. + - `static_4`: Select 4 for all blocks. + - `static_6`: Select 6 for all blocks (normal NVFP4 quantization). + - `mae`: Between 4 and 6, select the block scale that minimizes the mean absolute + quantization error. + - `mse`: Between 4 and 6, select the block scale that minimizes the mean squared + quantization error. + """ + + abs_max = "abs_max" + mae = "mae" + mse = "mse" + static_4 = "static_4" + static_6 = "static_6" + + def cuda_id(self) -> int: + """ID for the rule in the CUDA implementation.""" + + return { + ScaleRule.abs_max: 4, + ScaleRule.mae: 2, + ScaleRule.mse: 3, + ScaleRule.static_4: 1, + ScaleRule.static_6: 0, + }[self] + + def is_static(self) -> bool: + """Return True if the rule is static, False otherwise.""" + return self in {ScaleRule.static_4, ScaleRule.static_6} + + def max_allowed_e2m1_value(self) -> int: + """Return the maximum allowed E2M1 value for the rule.""" + return 4 if self == ScaleRule.static_4 else 6 + + def max_allowed_e4m3_value(self) -> int: + """Return the maximum allowed E4M3 value for the rule.""" + return 448 if self in {ScaleRule.static_6, ScaleRule.static_4} else 256 diff --git a/fouroversix/src/fouroversix/weight_conversions/__init__.py b/fouroversix/src/fouroversix/weight_conversions/__init__.py new file mode 100644 index 0000000..070f516 --- /dev/null +++ b/fouroversix/src/fouroversix/weight_conversions/__init__.py @@ -0,0 +1,13 @@ +import warnings + +try: + from .conversions import WeightConversions + from .gpt_oss import FourOverSixGptOssDeserialize, GptOssWeightConverter +except ImportError: + warnings.warn("Install transformers>=5.0 to use weight conversions", stacklevel=2) + + WeightConversions = None + FourOverSixGptOssDeserialize = None + GptOssWeightConverter = None + +__all__ = ["FourOverSixGptOssDeserialize", "GptOssWeightConverter", "WeightConversions"] diff --git a/fouroversix/src/fouroversix/weight_conversions/conversions.py b/fouroversix/src/fouroversix/weight_conversions/conversions.py new file mode 100644 index 0000000..f597e2d --- /dev/null +++ b/fouroversix/src/fouroversix/weight_conversions/conversions.py @@ -0,0 +1,45 @@ +from collections.abc import Callable +from typing import ClassVar + +from transformers import WeightConverter + + +class WeightConversions: + """Base class for weight conversions for quantized models.""" + + _registry: ClassVar[dict[str, list[WeightConverter]]] = {} + + @classmethod + def register( + cls, + pre_quantized_model_config_type: str, + ) -> Callable[[type], list[WeightConverter]]: + """Register a new type of weight conversion.""" + if pre_quantized_model_config_type in cls._registry: + msg = f"Model with config {pre_quantized_model_config_type} is \ + already registered." + raise ValueError(msg) + + def inner_wrapper( + wrapped_cls: type, + ) -> list[WeightConverter]: + weight_conversions = wrapped_cls.get_weight_conversions() + cls._registry[pre_quantized_model_config_type] = weight_conversions + return weight_conversions + + return inner_wrapper + + @classmethod + def get_weight_conversions( + cls, + pre_quantized_model_config_type: str, + ) -> list[WeightConverter]: + """ + Get the weight conversion for a given model type determined + by the model config type. + """ + weight_conversions = cls._registry.get(pre_quantized_model_config_type, None) + if weight_conversions is None: + msg = f"Config type {pre_quantized_model_config_type} not supported." + raise ValueError(msg) + return weight_conversions diff --git a/fouroversix/src/fouroversix/weight_conversions/gpt_oss.py b/fouroversix/src/fouroversix/weight_conversions/gpt_oss.py new file mode 100644 index 0000000..3266095 --- /dev/null +++ b/fouroversix/src/fouroversix/weight_conversions/gpt_oss.py @@ -0,0 +1,99 @@ +import torch +from fouroversix import DataType, ScaleRule +from fouroversix.quantize import QuantizedTensor +from transformers import ConversionOps, GptOssConfig, WeightConverter + +from .conversions import WeightConversions + + +class FourOverSixGptOssDeserialize(ConversionOps): + """Fouroversix deserializer for gpt oss model.""" + + def __init__( + self, + dtype: DataType = None, + scale_rule: ScaleRule = None, + ) -> None: + self.dtype = dtype + self.scale_rule = scale_rule + + def convert( + self, + input_dict: torch.Tensor, + **kwargs, # noqa: ARG002, ANN003 + ) -> dict[str, list[torch.Tensor]]: + """ + Convert the quantized parameters in gpt oss model to + high precision weights. + """ + + prefix = "" + if ".down_proj_blocks" in input_dict: + weight = input_dict[".down_proj_blocks"][0] + scales = input_dict[".down_proj_scales"][0] + prefix = "down" + elif ".gate_up_proj_blocks" in input_dict: + weight = input_dict[".gate_up_proj_blocks"][0] + scales = input_dict[".gate_up_proj_scales"][0] + prefix = "gate_up" + + num_experts = weight.shape[0] + hidden_size = weight.shape[1] + weight = weight.reshape((num_experts, hidden_size, -1)) + + dequantized_proj = [] + for e in range(num_experts): + weight_uint8 = weight[e].to(torch.uint8) + quantized_tensor = QuantizedTensor( + values=weight_uint8, + scale_factors=scales[e].to(torch.uint8).view(self.dtype.scale_dtype()), + amax=torch.ones( + (1,), + device=weight[e].device, + dtype=torch.float32, + ), + dtype=self.dtype, + original_shape=( + weight[e].shape[0], + weight[e].shape[1] * 2, + ), + scale_rule=self.scale_rule, + ) + + dequantized = quantized_tensor.dequantize() + dequantized_proj.append(dequantized) + + dequantized_weight = torch.stack(dequantized_proj, dim=0) + + return {f"{prefix}_proj": [dequantized_weight]} + + +@WeightConversions.register(str(GptOssConfig)) +class GptOssWeightConverter: + """Stores the weight conversions for the gpt oss model.""" + + @classmethod + def get_weight_conversions(cls) -> list[WeightConverter]: + """Return weight conversions for the gpt oss model.""" + return [ + WeightConverter( + source_patterns=[".gate_up_proj_blocks", ".gate_up_proj_scales"], + target_patterns=".gate_up_proj", + operations=[ + FourOverSixGptOssDeserialize( + dtype=DataType.mxfp4, + scale_rule=ScaleRule.static_6, + ), + ], + ), + WeightConverter( + source_patterns=[".down_proj_blocks", ".down_proj_scales"], + target_patterns=".down_proj", + operations=[ + FourOverSixGptOssDeserialize( + dtype=DataType.mxfp4, + scale_rule=ScaleRule.static_6, + ), + ], + ), + ] diff --git a/fouroversix/tests/__init__.py b/fouroversix/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/fouroversix/tests/test_correctness.py b/fouroversix/tests/test_correctness.py new file mode 100644 index 0000000..3fad9f5 --- /dev/null +++ b/fouroversix/tests/test_correctness.py @@ -0,0 +1,196 @@ +import itertools + +import pytest +import torch +from fouroversix import ( + DataType, + QuantizationConfig, + QuantizeBackend, + RoundStyle, + ScaleRule, + quantize_to_fp4, +) +from fouroversix.quantize.frontend import AVAILABLE_BACKENDS +from fouroversix.quantize.quantized_tensor import from_blocked + +MAE_MSE_MISMATCH_TOLERANCE = 1e-3 +NUM_RANDOM_SEEDS = 10 + + +@pytest.mark.parametrize("input_type", ["zeros", "ones", "rand01", "randn"]) +@pytest.mark.parametrize( + "input_shape", + [(1024, 1024), (1024, 512), (512, 1024)], +) +@pytest.mark.parametrize( + ("backend_a", "backend_b"), + itertools.combinations( + [ + QuantizeBackend.cuda, + QuantizeBackend.triton, + QuantizeBackend.pytorch, + QuantizeBackend.transformer_engine, + ], + r=2, + ), +) +@pytest.mark.parametrize("block_scale_2d", ["block_scale_2d", "no_block_scale_2d"]) +@pytest.mark.parametrize("dtype", [DataType.nvfp4]) +@pytest.mark.parametrize("rht", ["rht", "no_rht"]) +@pytest.mark.parametrize( + "scale_rule", + [ + ScaleRule.abs_max, + ScaleRule.mae, + ScaleRule.mse, + ScaleRule.static_4, + ScaleRule.static_6, + ], +) +@pytest.mark.parametrize("round_style", [RoundStyle.nearest, RoundStyle.stochastic]) +@pytest.mark.parametrize("transpose", ["transpose", "no_transpose"]) +def test_backend_outputs_are_consistent( # noqa: C901, PLR0915 + input_type: str, + input_shape: tuple[int, int], + backend_a: QuantizeBackend, + backend_b: QuantizeBackend, + *, + block_scale_2d: str, + dtype: DataType, + rht: str, + round_style: RoundStyle, + scale_rule: ScaleRule, + transpose: str, +) -> None: + torch.set_printoptions(precision=10) + + backend_a_cls = AVAILABLE_BACKENDS[backend_a] + backend_b_cls = AVAILABLE_BACKENDS[backend_b] + + if not backend_a_cls.is_available() or not backend_b_cls.is_available(): + pytest.skip("Backend is not available") + + config_a = QuantizationConfig( + backend=backend_a, + block_scale_2d=block_scale_2d == "block_scale_2d", + dtype=dtype, + rht=rht == "rht", + round_style=round_style, + scale_rule=scale_rule, + transpose=transpose == "transpose", + ) + + config_b = QuantizationConfig( + backend=backend_b, + block_scale_2d=block_scale_2d == "block_scale_2d", + dtype=dtype, + rht=rht == "rht", + round_style=round_style, + scale_rule=scale_rule, + transpose=transpose == "transpose", + ) + + if round_style == RoundStyle.stochastic: + pytest.xfail("This test is not currently targeting stochastic rounding") + + for random_seed in range(NUM_RANDOM_SEEDS): + print(f"Testing with random seed: {random_seed}") + torch.manual_seed(random_seed) + + if input_type == "zeros": + x = torch.zeros(*input_shape, dtype=torch.bfloat16, device="cuda") + elif input_type == "ones": + x = torch.ones(*input_shape, dtype=torch.bfloat16, device="cuda") + elif input_type == "rand01": + x = torch.randint(0, 2, input_shape, dtype=int, device="cuda").to( + torch.bfloat16, + ) + elif input_type == "randn": + x = torch.randn(*input_shape, dtype=torch.bfloat16, device="cuda") + else: + msg = f"Invalid input type: {input_type}" + raise ValueError(msg) + + if not backend_a_cls.is_supported( + x, + config_a, + ) or not backend_b_cls.is_supported( + x, + config_b, + ): + pytest.skip("Backend is not supported") + + quantized_a = quantize_to_fp4(x.clone(), config_a) + quantized_b = quantize_to_fp4(x.clone(), config_b) + + if not torch.allclose(quantized_a.amax, quantized_b.amax): + print("Backends A and B have different amax values!") + print(f"{backend_a}: {quantized_a.amax}") + print(f"{backend_b}: {quantized_b.amax}") + pytest.fail("Backends A and B have different amax values!") + + sf_a = from_blocked( + quantized_a.scale_factors.bfloat16(), + (input_shape[0], input_shape[1] // 16), + ) + sf_b = from_blocked( + quantized_b.scale_factors.bfloat16(), + (input_shape[0], input_shape[1] // 16), + ) + + # When computing 4/6 with the MAE and MSE scale rules, computing the errors + # requires summing the errors in each block of 16 values. This operation + # differently (elements are summed in different orders, and floating-point + # addition is not associative) in PyTorch and Triton, and can not be easily made + # deterministic in a way that allows for good performance. As a result, we allow + # a small number of mismatches between the scale factors and values for these + # two rules. Fortunately, abs_max does not involve a summation, so we can use it + # to test the correctness of the rest of the 4/6 implementation. + scale_factors_mismatch_prop = (sf_a != sf_b).sum() / sf_a.numel() + + if ( + scale_rule in {ScaleRule.static_6, ScaleRule.static_4, ScaleRule.abs_max} + and scale_factors_mismatch_prop > 0 + ) or scale_factors_mismatch_prop >= MAE_MSE_MISMATCH_TOLERANCE: + print( + "Backends A and B have different scale factors! " + f"{scale_factors_mismatch_prop:.2%} mismatch", + ) + + [i, *_], [j, *_] = torch.where(sf_a != sf_b) + print(backend_a) + print("sf", sf_a[i, j]) + print("e2m1", quantized_a.values[i, 8 * j : 8 * (j + 1)]) + print(backend_b) + print("sf", sf_b[i, j]) + print("e2m1", quantized_b.values[i, 8 * j : 8 * (j + 1)]) + print("original") + print("x", x[i, 16 * j : 16 * (j + 1)]) + pytest.fail("Backends A and B have different scale factors!") + + values_mismatch_prop = ( + quantized_a.values != quantized_b.values + ).sum() / quantized_a.values.numel() + + if ( + scale_rule in {ScaleRule.static_6, ScaleRule.static_4, ScaleRule.abs_max} + and values_mismatch_prop > 0 + ) or values_mismatch_prop >= MAE_MSE_MISMATCH_TOLERANCE: + print( + "Backends A and B have different e2m1 values! " + f"{values_mismatch_prop:.2%} mismatch", + ) + + [i, *_], [j, *_] = torch.where( + quantized_a.values != quantized_b.values, + ) + print(i, j) + print("amax", quantized_a.amax) + print("sf", sf_a[i, j // 8]) + print(backend_a) + print("e2m1", quantized_a.values[i, 8 * (j // 8) : 8 * (j // 8 + 1)]) + print(backend_b) + print("e2m1", quantized_b.values[i, 8 * (j // 8) : 8 * (j // 8 + 1)]) + print("original") + print("x", x[i, 16 * (j // 8) : 16 * (j // 8 + 1)]) + pytest.fail("Backends A and B have different e2m1 values!") diff --git a/inference.py b/inference.py new file mode 100644 index 0000000..58baaa1 --- /dev/null +++ b/inference.py @@ -0,0 +1,670 @@ +# Adopted from https://github.com/guandeh17/Self-Forcing +# SPDX-License-Identifier: Apache-2.0 +import os +import sys + +# torchrun no longer consistently prepends the script directory to sys.path, +# which breaks absolute project imports when launched from another cwd. +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +# torchvision 0.27+ removed write_video/read_video. Several modules import the +# symbols at module import time, so patch them before importing project code. +import torchvision.io as _tv_io +if not hasattr(_tv_io, "write_video"): + import imageio.v2 as _imageio_v2 + + def _shim_write_video(filename, video_array, fps, **_unused): + if hasattr(video_array, "detach"): + video_array = video_array.detach().cpu().numpy() + _imageio_v2.mimwrite(filename, video_array, fps=fps, codec="libx264", quality=8) + + _tv_io.write_video = _shim_write_video +if not hasattr(_tv_io, "read_video"): + import imageio.v3 as _imageio_v3 + import torch as _torch_for_shim + + def _shim_read_video(filename, pts_unit="sec", output_format="THWC", **_unused): + frames = _imageio_v3.imread(filename, plugin="pyav") + tensor = _torch_for_shim.from_numpy(frames) + if output_format == "TCHW": + tensor = tensor.permute(0, 3, 1, 2) + return tensor, None, {} + + _tv_io.read_video = _shim_read_video + +import argparse +import torch +from omegaconf import OmegaConf +from tqdm import tqdm +from torchvision.io import write_video +from einops import rearrange +import torch.distributed as dist +from torch.utils.data import DataLoader, SequentialSampler +from torch.utils.data.distributed import DistributedSampler + +from pipeline import CausalDiffusionInferencePipeline +from utils.dataset import MultiTextConcatDataset, MultiVideoConcatDataset, eval_collate_fn, multi_video_collate_fn +from utils.misc import set_seed +from utils.config import normalize_config, section_get, wan_default_config +from utils.nvfp4_checkpoint import ( + clean_fsdp_state_dict_keys, + drop_fouroversix_master_weights, + is_nvfp4_state_dict, + is_te_nvfp4_checkpoint, + quantize_model_for_fouroversix_nvfp4, + unwrap_generator_state_dict, +) + +from utils.memory import get_cuda_free_memory_gb, DynamicSwapInstaller + + +def save_prompts_to_txt(prompts_for_sample, prompt_txt_path: str, is_main_process: bool): + """Save per-block prompts alongside the video. + + Consecutive identical prompts are merged, e.g.: + [0] a, [1] a, [2] b => [0,1] a\\n[2] b\\n + """ + try: + with open(prompt_txt_path, "w", encoding="utf-8") as f: + if len(prompts_for_sample) == 0: + return + current_prompt = prompts_for_sample[0] + current_indices = [0] + for seg_idx in range(1, len(prompts_for_sample)): + p = prompts_for_sample[seg_idx] + if p == current_prompt: + current_indices.append(seg_idx) + else: + indices_str = ",".join(str(i) for i in current_indices) + f.write(f"[{indices_str}] {current_prompt}\n") + current_prompt = p + current_indices = [seg_idx] + indices_str = ",".join(str(i) for i in current_indices) + f.write(f"[{indices_str}] {current_prompt}\n") + except Exception as e: + if is_main_process: + print(f"Warning: failed to save prompts to {prompt_txt_path}: {e}") + +parser = argparse.ArgumentParser() +parser.add_argument("--config_path", type=str, help="Path to the config file") +te_quant_group = parser.add_mutually_exclusive_group() +te_quant_group.add_argument( + "--use_te_quant", + dest="use_te_quant", + action="store_true", + help="Override config and enable TransformerEngine quantization", +) +te_quant_group.add_argument( + "--no_use_te_quant", + dest="use_te_quant", + action="store_false", + help="Override config and disable TransformerEngine quantization", +) +parser.set_defaults(use_te_quant=None) +args = parser.parse_args() + +config = normalize_config(OmegaConf.load(args.config_path)) +if args.use_te_quant is not None: + config.model_quant_use_transformer_engine = args.use_te_quant + +if not hasattr(config, "sampling_steps") or config.sampling_steps is None: + raise ValueError("sampling_steps must be defined in the inference config") + +if not hasattr(config, "guidance_scale") or config.guidance_scale is None: + config.guidance_scale = 1.0 + +config.use_ema = section_get(config, "inference", "use_ema", getattr(config, "use_ema", False)) +config.output_folder = section_get(config, "inference", "output_folder", getattr(config, "output_folder", "videos/longlive2")) +config.num_samples = section_get(config, "inference", "num_samples", getattr(config, "num_samples", 1)) +config.num_output_frames = getattr(config, "num_output_frames", config.image_or_video_shape[1]) +config.save_with_index = getattr(config, "save_with_index", False) +config.inference_iter = getattr(config, "inference_iter", -1) + +if bool(getattr(config, "fp8_quant", False)) and bool( + getattr(config, "model_quant", False) +): + raise ValueError("fp8_quant and model_quant (NVFP4) are mutually exclusive.") + + +def _maybe_to_dict(value): + if value is None: + return None + if OmegaConf.is_config(value): + value = OmegaConf.to_container(value, resolve=True) + return dict(value) + + +def _config_bool(value, default=False): + if value is None: + return default + if isinstance(value, str): + return value.strip().lower() in {"1", "true", "yes", "y", "on"} + return bool(value) + + +def _expected_inference_samples(config): + inference_iter = int(getattr(config, "inference_iter", -1)) + if inference_iter >= 0: + return inference_iter + 1 + return None + + +def _resolve_torch_compile(config): + setting = getattr(config, "torch_compile", False) + if isinstance(setting, str) and setting.strip().lower() == "auto": + if not ( + bool(getattr(config, "model_quant", False)) + or bool(getattr(config, "fp8_quant", False)) + ): + return False, "auto disabled because quantization is false" + min_samples = int(getattr(config, "torch_compile_min_samples", 2)) + expected_samples = _expected_inference_samples(config) + if expected_samples is not None and expected_samples < min_samples: + return ( + False, + "auto disabled because expected samples " + f"({expected_samples}) < torch_compile_min_samples ({min_samples})", + ) + return True, "auto enabled for repeated quantized inference" + return _config_bool(setting, default=False), "explicit setting" + + +def quantize_generator_model(model, config, keep_master_weights): + from utils.quant import ( + ModelQuantizationConfig, + _materialize_mixed_quantized_weights_for_inference, + _materialize_quantized_weights_for_inference, + _materialize_transformer_engine_weights_for_inference, + quantize_model_with_filter, + ) + + use_transformer_engine = bool(getattr(config, "model_quant_use_transformer_engine", False)) + te_inference_only = bool(getattr(config, "model_quant_te_inference_only", use_transformer_engine)) + te_low_precision_weights = bool(getattr(config, "model_quant_te_low_precision_weights", te_inference_only)) + te_fallback_to_fouroversix = bool(getattr(config, "model_quant_te_fallback_to_fouroversix", False)) + + quant_cfg = ModelQuantizationConfig( + scale_rule=getattr(config, "model_quant_scale_rule", "static_6"), + quantize_backend=getattr(config, "model_quant_backend", None), + activation_scale_rule=getattr( + config, + "model_quant_activation_scale_rule", + getattr(config, "model_quant_scale_rule", "static_6"), + ), + weight_scale_rule=getattr(config, "model_quant_weight_scale_rule", None), + gradient_scale_rule=getattr(config, "model_quant_gradient_scale_rule", None), + ) + quant_cfg.keep_master_weights = keep_master_weights + model, matched_modules = quantize_model_with_filter( + model, + quant_config=quant_cfg, + filtered_modules=getattr(config, "model_quant_filtered_modules", None), + use_default_filtered_modules=getattr(config, "model_quant_use_default_filtered_modules", True), + cast_model_to_bf16=True, + materialize_for_inference=False, + use_transformer_engine=use_transformer_engine, + te_inference_only=te_inference_only, + te_low_precision_weights=te_low_precision_weights, + te_recipe_kwargs=_maybe_to_dict(getattr(config, "model_quant_te_recipe_kwargs", None)), + te_module_kwargs=_maybe_to_dict(getattr(config, "model_quant_te_module_kwargs", None)), + te_fallback_to_fouroversix=te_fallback_to_fouroversix, + verbose=True, + ) + materialize_fn = _materialize_quantized_weights_for_inference + if use_transformer_engine and te_fallback_to_fouroversix: + materialize_fn = _materialize_mixed_quantized_weights_for_inference + elif use_transformer_engine: + materialize_fn = _materialize_transformer_engine_weights_for_inference + if local_rank == 0: + print(f"[NVFP4] Generator quantized; {len(matched_modules)} modules excluded") + return model, materialize_fn + + +def materialize_quantized_generator(model, device, materialize_fn, stage_desc): + mat_modules, master_bytes, quantized_bytes = materialize_fn( + model, + target_device=device, + ) + if local_rank == 0: + print( + f"[NVFP4] Materialized quantized generator weights {stage_desc}: " + f"{len(mat_modules)} modules, " + f"master_weight={master_bytes / (1024 ** 3):.3f} GiB, " + f"quantized_weight={quantized_bytes / (1024 ** 3):.3f} GiB" + ) + + +def configure_generator_torch_compile(pipeline, config): + compile_enabled, reason = _resolve_torch_compile(config) + if not compile_enabled: + if local_rank == 0 and str(getattr(config, "torch_compile", "false")).lower() == "auto": + print(f"[torch.compile] skipped: {reason}") + return + target = str(getattr(config, "torch_compile_target", "generator_model")).lower() + if target not in {"generator_model", "model"}: + if local_rank == 0: + print(f"[torch.compile][warn] Unsupported target={target}; expected generator_model") + return + if not hasattr(pipeline.generator, "configure_torch_compile"): + if local_rank == 0: + print("[torch.compile][warn] Current generator does not expose configure_torch_compile; skipping") + return + compiled = pipeline.generator.configure_torch_compile( + backend=str(getattr(config, "torch_compile_backend", "inductor")), + mode=getattr(config, "torch_compile_mode", "max-autotune-no-cudagraphs"), + fullgraph=_config_bool(getattr(config, "torch_compile_fullgraph", False)), + dynamic=_config_bool(getattr(config, "torch_compile_dynamic", False)), + options=_maybe_to_dict(getattr(config, "torch_compile_options", None)), + suppress_errors=_config_bool(getattr(config, "torch_compile_suppress_errors", True), default=True), + ) + if local_rank == 0: + status = "enabled" if compiled else "not enabled" + print(f"[torch.compile] {status}: target={target}") + +# Initialize distributed inference +if "LOCAL_RANK" in os.environ: + dist.init_process_group(backend='nccl') + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + set_seed(config.seed + local_rank) + config.distributed = True # Mark as distributed for pipeline +else: + local_rank = 0 + device = torch.device("cuda") + set_seed(config.seed) + config.distributed = False # Mark as non-distributed + +print(f'Free VRAM {get_cuda_free_memory_gb(device)} GB') +low_memory = get_cuda_free_memory_gb(device) < 40 + +torch.set_grad_enabled(False) + + +# Initialize pipeline +pipeline = CausalDiffusionInferencePipeline(config, device=device) + +# --------------------------- LoRA support (optional) --------------------------- +from utils.lora_utils import configure_lora_for_model +import peft + +merge_lora = bool(getattr(config, "merge_lora", False)) +has_lora_adapter = bool(getattr(config, "adapter", None) and configure_lora_for_model is not None) +if has_lora_adapter and ( + bool(getattr(config, "model_quant", False)) + or bool(getattr(config, "fp8_quant", False)) +) and not merge_lora: + if local_rank == 0: + print( + "[quant][LoRA] merge_lora=false is unsupported with quantization; " + "forcing merge_lora=true so the LoRA is folded into the BF16 base before quantization." + ) + merge_lora = True + config.merge_lora = True +materialize_quantized_weights_for_inference = None +generator_checkpoint = None +generator_lora_state = None +generator_ckpt_path = getattr(config, "generator_ckpt", None) +loaded_prequantized_generator = False +prequantized_generator_backend = None +if generator_ckpt_path: + generator_checkpoint = torch.load(generator_ckpt_path, map_location="cpu") + is_lora_only_checkpoint = ( + isinstance(generator_checkpoint, dict) + and "generator_lora" in generator_checkpoint + and not any(key in generator_checkpoint for key in ("generator", "generator_ema", "model")) + ) + if is_lora_only_checkpoint: + generator_lora_state = generator_checkpoint["generator_lora"] + if local_rank == 0: + print(f"Found LoRA generator weights in {generator_ckpt_path}") + else: + raw_gen_state_dict = unwrap_generator_state_dict(generator_checkpoint, use_ema=config.use_ema) + if config.use_ema: + raw_gen_state_dict = clean_fsdp_state_dict_keys(raw_gen_state_dict) + if is_te_nvfp4_checkpoint(generator_checkpoint): + raise ValueError( + "This checkpoint was saved as a TransformerEngine module state_dict, " + "which is not packed NVFP4 and is no longer a supported export format. " + "Regenerate with `--backend transformer_engine` to save merged bf16 weights " + "for TE runtime quantization, or use `--backend fouroversix` for a compact " + "materialized NVFP4 checkpoint." + ) + elif is_nvfp4_state_dict(raw_gen_state_dict): + if not getattr(config, "model_quant", False): + raise ValueError( + "generator_ckpt is a materialized NVFP4 checkpoint, but model_quant is false. " + "Set model_quant: true in the inference yaml." + ) + if getattr(config, "model_quant_use_transformer_engine", False): + raise ValueError( + "Materialized NVFP4 generator checkpoints use FourOverSix modules. " + "Set model_quant_use_transformer_engine: false when loading this checkpoint." + ) + if local_rank == 0: + print(f"[NVFP4] Loading materialized generator checkpoint from {generator_ckpt_path}") + pipeline.generator.model, matched_modules = quantize_model_for_fouroversix_nvfp4( + pipeline.generator.model, + config=config, + keep_master_weights=False, + verbose=(local_rank == 0), + ) + dropped_modules = drop_fouroversix_master_weights(pipeline.generator.model) + pipeline.generator.load_state_dict(raw_gen_state_dict, strict=True) + loaded_prequantized_generator = True + prequantized_generator_backend = "fouroversix" + if local_rank == 0: + print( + "[NVFP4] Prepared quantized generator architecture: " + f"{len(dropped_modules)} materialized modules, " + f"{len(matched_modules)} modules excluded" + ) + elif config.use_ema: + missing, unexpected = pipeline.generator.load_state_dict(raw_gen_state_dict, strict=False) + if local_rank == 0: + if len(missing) > 0: + print(f"[Warning] {len(missing)} parameters are missing when loading checkpoint: {missing[:8]} ...") + if len(unexpected) > 0: + print(f"[Warning] {len(unexpected)} unexpected parameters encountered when loading checkpoint: {unexpected[:8]} ...") + else: + print(f"Loading generator from {generator_ckpt_path}") + pipeline.generator.load_state_dict(raw_gen_state_dict, strict=True) + +pipeline.is_lora_enabled = False +pipeline.is_lora_merged = False + +if loaded_prequantized_generator: + if has_lora_adapter or merge_lora or getattr(config, "lora_ckpt", None): + if local_rank == 0: + print("[NVFP4] Pre-quantized generator checkpoint is already saved with merged weights; skipping LoRA setup") + has_lora_adapter = False + merge_lora = False + config.merge_lora = False + +if getattr(config, "model_quant", False) and not merge_lora and not loaded_prequantized_generator: + pipeline.generator.model, materialize_quantized_weights_for_inference = quantize_generator_model( + pipeline.generator.model, + config=config, + keep_master_weights=has_lora_adapter, + ) + +if has_lora_adapter: + if local_rank == 0: + print(f"LoRA enabled with config: {config.adapter}") + print("Applying LoRA to generator (inference)...") + if merge_lora: + print("LoRA weights will be merged into the base model before inference") + # Apply LoRA to the generator transformer after loading base weights. + pipeline.generator.model = configure_lora_for_model( + pipeline.generator.model, + model_name="generator", + lora_config=config.adapter, + is_main_process=(local_rank == 0), + ) + + # Load LoRA weights from lora_ckpt. If omitted, fall back to generator_ckpt + # when it directly contains generator_lora. + lora_ckpt_path = getattr(config, "lora_ckpt", None) + if lora_ckpt_path: + if local_rank == 0: + print(f"Loading LoRA weights from lora_ckpt: {lora_ckpt_path}") + lora_checkpoint = torch.load(lora_ckpt_path, map_location="cpu") + if isinstance(lora_checkpoint, dict) and "generator_lora" in lora_checkpoint: + peft.set_peft_model_state_dict(pipeline.generator.model, lora_checkpoint["generator_lora"]) # type: ignore + else: + peft.set_peft_model_state_dict(pipeline.generator.model, lora_checkpoint) # type: ignore + if local_rank == 0: + print("LoRA weights loaded for generator") + elif generator_lora_state is not None: + if local_rank == 0: + print(f"Loading LoRA weights from generator_ckpt: {generator_ckpt_path}") + peft.set_peft_model_state_dict(pipeline.generator.model, generator_lora_state) # type: ignore + if local_rank == 0: + print("LoRA weights loaded for generator") + else: + if local_rank == 0: + print("No LoRA checkpoint configured; using initialized LoRA adapters") + + if merge_lora: + if local_rank == 0: + print("Merging LoRA weights into generator before quantization / inference...") + pipeline.generator.model = pipeline.generator.model.merge_and_unload(safe_merge=True) + pipeline.is_lora_merged = True + else: + pipeline.is_lora_enabled = True +elif merge_lora and local_rank == 0: + print("merge_lora=True requested but no adapter config was found; continuing without LoRA merge") + +del generator_checkpoint + + +# Move pipeline to appropriate dtype and device +if loaded_prequantized_generator: + pipeline.text_encoder.to(dtype=torch.bfloat16) + pipeline.vae.to(dtype=torch.bfloat16) +else: + pipeline = pipeline.to(dtype=torch.bfloat16) +if low_memory: + DynamicSwapInstaller.install_model(pipeline.text_encoder, device=device) +pipeline.generator.to(device=device) + +if getattr(config, "model_quant", False) and not loaded_prequantized_generator: + if merge_lora: + pipeline.generator.model, materialize_quantized_weights_for_inference = quantize_generator_model( + pipeline.generator.model, + config=config, + keep_master_weights=False, + ) + stage_desc = "after LoRA merge" if pipeline.is_lora_merged else "for inference" + else: + stage_desc = "after LoRA wrapping" if pipeline.is_lora_enabled else "for inference" + materialize_quantized_generator( + pipeline.generator.model, + device=device, + materialize_fn=materialize_quantized_weights_for_inference, + stage_desc=stage_desc, + ) +elif loaded_prequantized_generator and local_rank == 0: + print(f"[NVFP4] Using pre-saved {prequantized_generator_backend} generator weights from checkpoint") + +pipeline.generator.model.eval().requires_grad_(False) +if bool(getattr(config, "fp8_quant", False)): + from utils.fp8 import quantize_model_fp8 + + quantize_model_fp8(pipeline.generator.model, verbose=(local_rank == 0)) +configure_generator_torch_compile(pipeline, config) + +vae_device_str = getattr(config, "vae_device", None) +use_dedicated_vae_device = bool(getattr(config, "streaming_vae", False)) and bool(vae_device_str) +if use_dedicated_vae_device: + vae_device = torch.device(vae_device_str) + pipeline.vae.to(device="cpu") + pipeline.vae.to(device=vae_device) + if hasattr(pipeline.vae, "mean"): + pipeline.vae.mean = pipeline.vae.mean.to(device=vae_device) + pipeline.vae.std = pipeline.vae.std.to(device=vae_device) + if local_rank == 0: + print(f"[inference] VAE on {vae_device}, diffusion on {device}") +else: + pipeline.vae.to(device=device) + if vae_device_str and local_rank == 0: + print(f"[inference] Ignoring vae_device={vae_device_str} because streaming_vae is false") + +# Create dataset +nfpb = getattr(config, 'num_frame_per_block', 8) +data_path = config.data_path +chunks_per_shot = getattr(config, 'chunks_per_shot', 0) +scene_cut_prefix = getattr(config, 'scene_cut_prefix', "The scene transitions. ") +if getattr(config, "i2v", False): + model_name = config.model_kwargs.model_name + frame_raw_height = list(config.image_or_video_shape)[3] * wan_default_config[model_name]["spatial_compression_ratio"] + frame_raw_width = list(config.image_or_video_shape)[4] * wan_default_config[model_name]["spatial_compression_ratio"] + temporal_compression_ratio = wan_default_config[model_name]["temporal_compression_ratio"] + total_frames = (config.num_output_frames - 1) * temporal_compression_ratio + 1 + dataset = MultiVideoConcatDataset( + data_dir=data_path, + video_size=(frame_raw_height, frame_raw_width), + total_frames=total_frames, + deterministic=True, + num_frame_per_block=nfpb, + temporal_compression_ratio=temporal_compression_ratio, + target_fps=24 if "5B" in model_name else 16, + allow_padding=getattr(config, "allow_padding", False), + min_latent_frames=getattr(config, "min_latent_frames", 0), + single_video_only=getattr(config, "uniform_prompt", False), + independent_first_frame=getattr(config, "independent_first_frame", False), + return_image=True, + max_chunks_per_shot=getattr(config, "max_chunks_per_shot", 0), + scene_cut_prefix=scene_cut_prefix, + ) + collate_fn = multi_video_collate_fn + num_blocks = config.num_output_frames // nfpb +else: + num_blocks = config.num_output_frames // nfpb + dataset = MultiTextConcatDataset( + data_path=data_path, + num_blocks=num_blocks, + chunks_per_shot=chunks_per_shot, + scene_cut_prefix=scene_cut_prefix, + deterministic=True, + ) + collate_fn = eval_collate_fn +if local_rank == 0: + print(f"[data] data_path={data_path}, mode={getattr(dataset, '_mode', dataset.__class__.__name__)}, num_blocks={num_blocks}") +num_prompts = len(dataset) +print(f"Number of prompts: {num_prompts}") + +if dist.is_initialized(): + sampler = DistributedSampler(dataset, shuffle=False, drop_last=True) +else: + sampler = SequentialSampler(dataset) +dataloader = DataLoader(dataset, batch_size=1, sampler=sampler, num_workers=0, + drop_last=False, collate_fn=collate_fn) + +# Create output directory (only on main process to avoid race conditions) +if local_rank == 0: + os.makedirs(config.output_folder, exist_ok=True) + +if dist.is_initialized(): + dist.barrier() + + +def encode(self, videos: torch.Tensor) -> torch.Tensor: + device, dtype = videos[0].device, videos[0].dtype + scale = [self.mean.to(device=device, dtype=dtype), + 1.0 / self.std.to(device=device, dtype=dtype)] + output = [ + self.model.encode(u.unsqueeze(0), scale).float().squeeze(0) + for u in videos + ] + + output = torch.stack(output, dim=0) + return output + + +for i, batch_data in tqdm(enumerate(dataloader), disable=(local_rank != 0)): + idx = batch_data['idx'].item() + + # For DataLoader batch_size=1, the batch_data is already a single item, but in a batch container + # Unpack the batch data for convenience + if isinstance(batch_data, dict): + batch = batch_data + elif isinstance(batch_data, list): + batch = batch_data[0] # First (and only) item in the batch + + all_video = [] + + # MultiTextConcatDataset + eval_collate_fn: prompts[0] is List[str]. + block_prompts = list(batch['prompts'][0]) + prompt = block_prompts[0] # for filename + prompts = [block_prompts] * config.num_samples + + shape = config.image_or_video_shape + sampled_noise = torch.randn( + [config.num_samples, config.num_output_frames, shape[2], shape[3], shape[4]], device=device, dtype=torch.bfloat16 + ) + initial_latent = None + if getattr(config, "i2v", False): + image = batch["image"].to(device=device, dtype=torch.bfloat16) + if image.ndim == 4: + image = image.unsqueeze(2) + elif image.ndim != 5: + raise ValueError(f"Expected i2v image with shape [B,C,H,W] or [B,C,T,H,W], got {tuple(image.shape)}") + initial_latent = pipeline.vae.encode_to_latent(image).to(device=device, dtype=torch.bfloat16) + if initial_latent.shape[0] != config.num_samples: + initial_latent = initial_latent.repeat(config.num_samples, 1, 1, 1, 1) + if config.num_output_frames <= initial_latent.shape[1]: + raise ValueError( + f"num_output_frames must exceed the i2v conditioning frames; " + f"got {config.num_output_frames} and {initial_latent.shape[1]}" + ) + print("sampled_noise.device", sampled_noise.device) + print("prompts", prompts) + print('sampled_noise.shape', sampled_noise.shape, 'prompts', prompts) + save_latents_only = section_get( + config, + "inference", + "save_latents_only", + getattr(config, "save_latents_only", getattr(config, "save_latent_only", False)), + aliases=("save_latent_only", "return_latents"), + ) + inference_kwargs = dict( + noise=sampled_noise, + text_prompts=prompts, + return_latents=save_latents_only, + ) + if initial_latent is not None: + inference_kwargs["initial_latent"] = initial_latent + with torch.inference_mode(): + generated = pipeline.inference(**inference_kwargs) + + if not save_latents_only: + current_video = rearrange(generated, 'b t c h w -> b t h w c').cpu() + all_video.append(current_video) + + # Final output video + video = 255.0 * torch.cat(all_video, dim=1) + + # Clear VAE cache + pipeline.vae.model.clear_cache() + else: + latents = generated + + if dist.is_initialized(): + rank = dist.get_rank() + else: + rank = 0 + + # Save the video if the current prompt is not a dummy prompt + if idx < num_prompts: + # Determine model type for filename + if hasattr(pipeline, 'is_lora_enabled') and pipeline.is_lora_enabled: + model_type = "lora" + elif getattr(config, 'use_ema', False): + model_type = "ema" + else: + model_type = "regular" + + for seed_idx in range(config.num_samples): + if config.save_with_index: + base_name = f'rank{rank}-{idx}-{seed_idx}_{model_type}' + else: + base_name = f'rank{rank}-{prompt[:100]}-{seed_idx}_{model_type}' + + if save_latents_only: + latent_path = os.path.join(config.output_folder, f'{base_name}.pt') + torch.save(latents[seed_idx].cpu(), latent_path) + else: + output_path = os.path.join(config.output_folder, f'{base_name}.mp4') + fps = 24 if '5B' in config.model_kwargs.model_name else 16 + write_video(output_path, video[seed_idx], fps=fps) + + prompt_txt_path = os.path.join(config.output_folder, f'{base_name}_prompts.txt') + save_prompts_to_txt( + prompts[seed_idx] if isinstance(prompts[seed_idx], list) else [prompts[seed_idx]], + prompt_txt_path, + is_main_process=(rank == 0), + ) + + if config.inference_iter != -1 and i >= config.inference_iter: + break diff --git a/inference_sp.py b/inference_sp.py new file mode 100644 index 0000000..b19b5ce --- /dev/null +++ b/inference_sp.py @@ -0,0 +1,578 @@ +# Adopted from https://github.com/guandeh17/Self-Forcing +# SPDX-License-Identifier: Apache-2.0 +import argparse +import os +from math import gcd + +import peft +import torch +import torch.distributed as dist +from einops import rearrange +from omegaconf import OmegaConf +from torch.utils.data import DataLoader, SequentialSampler +from torch.utils.data.distributed import DistributedSampler +from torchvision.io import write_video +from tqdm import tqdm + +from pipeline.causal_diffusion_inference_sp import CausalDiffusionInferencePipelineSP +from utils.config import normalize_config, section_get +from utils.dataset import MultiTextConcatDataset, eval_collate_fn +from utils.lora_utils import configure_lora_for_model +from utils.memory import DynamicSwapInstaller, get_cuda_free_memory_gb +from utils.misc import set_seed +from utils.nvfp4_checkpoint import ( + clean_fsdp_state_dict_keys, + drop_fouroversix_master_weights, + is_nvfp4_state_dict, + is_te_nvfp4_checkpoint, + quantize_model_for_fouroversix_nvfp4, + unwrap_generator_state_dict, +) + + +def save_prompts_to_txt(prompts_for_sample, prompt_txt_path: str, is_main_process: bool): + try: + with open(prompt_txt_path, "w", encoding="utf-8") as f: + if len(prompts_for_sample) == 0: + return + current_prompt = prompts_for_sample[0] + current_indices = [0] + for seg_idx in range(1, len(prompts_for_sample)): + prompt = prompts_for_sample[seg_idx] + if prompt == current_prompt: + current_indices.append(seg_idx) + else: + f.write(f"[{','.join(str(i) for i in current_indices)}] {current_prompt}\n") + current_prompt = prompt + current_indices = [seg_idx] + f.write(f"[{','.join(str(i) for i in current_indices)}] {current_prompt}\n") + except Exception as exc: + if is_main_process: + print(f"Warning: failed to save prompts to {prompt_txt_path}: {exc}") + + +def compute_group_specs(world_size, sp_size, dp_size, num_heads, num_frame_per_block, + auto_sp_remainder=False): + """Compute DP groups whose ranks each form a Ulysses SP group.""" + valid_base = gcd(num_heads, num_frame_per_block) + valid_sp_sizes = sorted(size for size in range(1, valid_base + 1) if valid_base % size == 0) + if sp_size not in valid_sp_sizes: + raise ValueError( + f"sp_size={sp_size} must divide gcd(num_heads={num_heads}, " + f"num_frame_per_block={num_frame_per_block})={valid_base}" + ) + + full_groups = min(world_size // max(sp_size, 1), dp_size) + groups = [ + (sp_size, list(range(i * sp_size, (i + 1) * sp_size))) + for i in range(full_groups) + ] + ranks_used = full_groups * sp_size + if auto_sp_remainder: + remaining_gpus = world_size - ranks_used + remaining_quota = dp_size - full_groups + while remaining_gpus > 0 and remaining_quota > 0: + size = max((v for v in valid_sp_sizes if v <= remaining_gpus), default=1) + groups.append((size, list(range(ranks_used, ranks_used + size)))) + ranks_used += size + remaining_gpus -= size + remaining_quota -= 1 + return groups + + +def _maybe_to_dict(value): + if value is None: + return None + if OmegaConf.is_config(value): + value = OmegaConf.to_container(value, resolve=True) + return dict(value) + + +def _config_bool(value, default=False): + if value is None: + return default + if isinstance(value, str): + return value.strip().lower() in {"1", "true", "yes", "y", "on"} + return bool(value) + + +def _expected_inference_samples(config): + inference_iter = int(getattr(config, "inference_iter", -1)) + return inference_iter + 1 if inference_iter >= 0 else None + + +def _resolve_torch_compile(config): + setting = getattr(config, "torch_compile", False) + if isinstance(setting, str) and setting.strip().lower() == "auto": + if not bool(getattr(config, "model_quant", False)): + return False, "auto disabled because model_quant is false" + min_samples = int(getattr(config, "torch_compile_min_samples", 2)) + expected_samples = _expected_inference_samples(config) + if expected_samples is not None and expected_samples < min_samples: + return False, f"auto disabled because expected samples ({expected_samples}) < {min_samples}" + return True, "auto enabled for repeated quantized inference" + return _config_bool(setting, default=False), "explicit setting" + + +def quantize_generator_model(model, config, keep_master_weights, is_main_process): + from utils.quant import ( + ModelQuantizationConfig, + _materialize_mixed_quantized_weights_for_inference, + _materialize_quantized_weights_for_inference, + _materialize_transformer_engine_weights_for_inference, + quantize_model_with_filter, + ) + + use_transformer_engine = bool(getattr(config, "model_quant_use_transformer_engine", False)) + te_inference_only = bool(getattr(config, "model_quant_te_inference_only", use_transformer_engine)) + te_low_precision_weights = bool(getattr(config, "model_quant_te_low_precision_weights", te_inference_only)) + te_fallback_to_fouroversix = bool(getattr(config, "model_quant_te_fallback_to_fouroversix", False)) + quant_cfg = ModelQuantizationConfig( + scale_rule=getattr(config, "model_quant_scale_rule", "static_6"), + quantize_backend=getattr(config, "model_quant_backend", None), + activation_scale_rule=getattr( + config, "model_quant_activation_scale_rule", + getattr(config, "model_quant_scale_rule", "static_6"), + ), + weight_scale_rule=getattr(config, "model_quant_weight_scale_rule", None), + gradient_scale_rule=getattr(config, "model_quant_gradient_scale_rule", None), + ) + quant_cfg.keep_master_weights = keep_master_weights + model, matched_modules = quantize_model_with_filter( + model, + quant_config=quant_cfg, + filtered_modules=getattr(config, "model_quant_filtered_modules", None), + use_default_filtered_modules=getattr(config, "model_quant_use_default_filtered_modules", True), + cast_model_to_bf16=True, + materialize_for_inference=False, + use_transformer_engine=use_transformer_engine, + te_inference_only=te_inference_only, + te_low_precision_weights=te_low_precision_weights, + te_recipe_kwargs=_maybe_to_dict(getattr(config, "model_quant_te_recipe_kwargs", None)), + te_module_kwargs=_maybe_to_dict(getattr(config, "model_quant_te_module_kwargs", None)), + te_fallback_to_fouroversix=te_fallback_to_fouroversix, + verbose=is_main_process, + ) + materialize_fn = _materialize_quantized_weights_for_inference + if use_transformer_engine and te_fallback_to_fouroversix: + materialize_fn = _materialize_mixed_quantized_weights_for_inference + elif use_transformer_engine: + materialize_fn = _materialize_transformer_engine_weights_for_inference + if is_main_process: + print(f"[NVFP4] Generator quantized; {len(matched_modules)} modules excluded") + return model, materialize_fn + + +def materialize_quantized_generator(model, device, materialize_fn, stage_desc, is_main_process): + mat_modules, master_bytes, quantized_bytes = materialize_fn(model, target_device=device) + if is_main_process: + print( + f"[NVFP4] Materialized quantized generator weights {stage_desc}: " + f"{len(mat_modules)} modules, master_weight={master_bytes / (1024 ** 3):.3f} GiB, " + f"quantized_weight={quantized_bytes / (1024 ** 3):.3f} GiB" + ) + + +def configure_generator_torch_compile(pipeline, config, is_main_process): + compile_enabled, reason = _resolve_torch_compile(config) + if not compile_enabled: + if is_main_process and str(getattr(config, "torch_compile", "false")).lower() == "auto": + print(f"[torch.compile] skipped: {reason}") + return + if not hasattr(pipeline.generator, "configure_torch_compile"): + if is_main_process: + print("[torch.compile][warn] Current generator does not expose configure_torch_compile; skipping") + return + compiled = pipeline.generator.configure_torch_compile( + backend=str(getattr(config, "torch_compile_backend", "inductor")), + mode=getattr(config, "torch_compile_mode", "max-autotune-no-cudagraphs"), + fullgraph=_config_bool(getattr(config, "torch_compile_fullgraph", False)), + dynamic=_config_bool(getattr(config, "torch_compile_dynamic", False)), + options=_maybe_to_dict(getattr(config, "torch_compile_options", None)), + suppress_errors=_config_bool(getattr(config, "torch_compile_suppress_errors", True), default=True), + ) + if is_main_process: + print(f"[torch.compile] {'enabled' if compiled else 'not enabled'}: target=generator_model") + + +parser = argparse.ArgumentParser() +parser.add_argument("--config_path", type=str, required=True, help="Path to the config YAML file") +te_quant_group = parser.add_mutually_exclusive_group() +te_quant_group.add_argument("--use_te_quant", dest="use_te_quant", action="store_true") +te_quant_group.add_argument("--no_use_te_quant", dest="use_te_quant", action="store_false") +parser.set_defaults(use_te_quant=None) +args = parser.parse_args() + +config = normalize_config(OmegaConf.load(args.config_path)) +if args.use_te_quant is not None: + config.model_quant_use_transformer_engine = args.use_te_quant +if not hasattr(config, "sampling_steps") or config.sampling_steps is None: + raise ValueError("sampling_steps must be defined in the SP inference config") +if not hasattr(config, "guidance_scale") or config.guidance_scale is None: + config.guidance_scale = 1.0 + +config.use_ema = section_get(config, "inference", "use_ema", getattr(config, "use_ema", False)) +config.output_folder = section_get(config, "inference", "output_folder", getattr(config, "output_folder", "videos/longlive2_sp")) +config.num_samples = section_get(config, "inference", "num_samples", getattr(config, "num_samples", 1)) +config.num_output_frames = getattr(config, "num_output_frames", config.image_or_video_shape[1]) +config.save_with_index = getattr(config, "save_with_index", False) +config.inference_iter = getattr(config, "inference_iter", -1) +if bool(getattr(config, "fp8_quant", False)): + raise NotImplementedError("TorchAO FP8 PTQ is currently supported by inference.py only.") +if getattr(config, "i2v", False): + raise NotImplementedError("I2V inference is not included in this SP release path.") +if getattr(config, "kv_quant", False): + print("[SP][warn] kv_quant is not supported in Ulysses SP inference; disabling it.") + config.kv_quant = False + +sp_size = int(getattr(config, "sp_size", 1)) +dp_size = int(getattr(config, "dp_size", 1)) +auto_sp_remainder = bool(getattr(config, "auto_sp_remainder", False)) +model_num_heads = int(getattr(config, "model_num_heads", 24)) + +sp_group = None +dp_rank = 0 +sp_rank = 0 +effective_sp_size = sp_size +total_dp_groups = 1 + +if "LOCAL_RANK" in os.environ: + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + rank = int(os.environ.get("RANK", str(local_rank))) + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + os.environ.setdefault("NCCL_CROSS_NIC", "1") + os.environ.setdefault("NCCL_DEBUG", "WARN") + os.environ.setdefault("NCCL_TIMEOUT", "1800") + if not dist.is_initialized(): + dist.init_process_group(backend="nccl", rank=rank, world_size=world_size) + set_seed(config.seed + rank) + config.distributed = True + + group_specs = compute_group_specs( + world_size=world_size, + sp_size=sp_size, + dp_size=dp_size, + num_heads=model_num_heads, + num_frame_per_block=int(getattr(config, "num_frame_per_block", 8)), + auto_sp_remainder=auto_sp_remainder, + ) + assigned_count = sum(len(ranks) for _, ranks in group_specs) + if assigned_count != world_size: + raise ValueError( + f"SP group layout assigns {assigned_count}/{world_size} ranks. " + "Increase dp_size, lower sp_size, or enable auto_sp_remainder." + ) + total_dp_groups = len(group_specs) + sp_groups_all = [dist.new_group(ranks=ranks) for _, ranks in group_specs] + for dp_i, (eff_sp, ranks) in enumerate(group_specs): + if rank in ranks: + dp_rank = dp_i + sp_rank = ranks.index(rank) + effective_sp_size = eff_sp + sp_group = sp_groups_all[dp_i] + break + if rank == 0: + print( + f"[SP] Parallelism: {total_dp_groups} DP group(s), " + f"sp_sizes={[size for size, _ in group_specs]}, assigned={assigned_count}/{world_size}" + ) +else: + local_rank = 0 + rank = 0 + device = torch.device("cuda") + set_seed(config.seed) + config.distributed = False + effective_sp_size = 1 + +is_main_process = rank == 0 +use_effective_sp = effective_sp_size > 1 and dist.is_initialized() +use_multi_dp = total_dp_groups > 1 + +if use_effective_sp: + from wan_5b.distributed.sp_ulysses_inference import init_sequence_parallel + init_sequence_parallel(group=sp_group) + if is_main_process: + print(f"[SP] Ulysses mode enabled: sp_sizes={[size for size, _ in group_specs]}") +elif is_main_process: + print("[SP] Running SP model with world_size=1") + +torch.set_grad_enabled(False) +free_vram = get_cuda_free_memory_gb(device) +low_memory = free_vram < 40 +if is_main_process: + print(f"[SP] Free VRAM: {free_vram:.1f} GB, low_memory={low_memory}") + +pipeline = CausalDiffusionInferencePipelineSP( + config, + device=device, + sp_group=sp_group, + dp_rank=dp_rank, +) + +merge_lora = bool(getattr(config, "merge_lora", False)) +has_lora_adapter = bool(getattr(config, "adapter", None) and configure_lora_for_model is not None) +if has_lora_adapter and bool(getattr(config, "model_quant", False)) and not merge_lora: + if is_main_process: + print( + "[NVFP4][LoRA] merge_lora=false is unsupported with model_quant=true; " + "forcing merge_lora=true so the LoRA is folded into the BF16 base before quantization." + ) + merge_lora = True + config.merge_lora = True +materialize_quantized_weights_for_inference = None +generator_checkpoint = None +generator_lora_state = None +generator_ckpt_path = getattr(config, "generator_ckpt", None) +loaded_prequantized_generator = False +prequantized_generator_backend = None + +if generator_ckpt_path: + if is_main_process: + print(f"[SP] Loading generator checkpoint: {generator_ckpt_path}") + generator_checkpoint = torch.load(generator_ckpt_path, map_location="cpu", mmap=True) + is_lora_only_checkpoint = ( + isinstance(generator_checkpoint, dict) + and "generator_lora" in generator_checkpoint + and not any(key in generator_checkpoint for key in ("generator", "generator_ema", "model")) + ) + if is_lora_only_checkpoint: + generator_lora_state = generator_checkpoint["generator_lora"] + else: + raw_gen_state_dict = unwrap_generator_state_dict(generator_checkpoint, use_ema=config.use_ema) + if config.use_ema: + raw_gen_state_dict = clean_fsdp_state_dict_keys(raw_gen_state_dict) + if is_te_nvfp4_checkpoint(generator_checkpoint): + raise ValueError("TransformerEngine module state_dict checkpoints are not supported here.") + if is_nvfp4_state_dict(raw_gen_state_dict): + if not getattr(config, "model_quant", False): + raise ValueError("generator_ckpt is materialized NVFP4 but model_quant is false.") + if getattr(config, "model_quant_use_transformer_engine", False): + raise ValueError("Materialized NVFP4 checkpoints require model_quant_use_transformer_engine=false.") + pipeline.generator.model, matched_modules = quantize_model_for_fouroversix_nvfp4( + pipeline.generator.model, + config=config, + keep_master_weights=False, + verbose=is_main_process, + ) + dropped_modules = drop_fouroversix_master_weights(pipeline.generator.model) + pipeline.generator.load_state_dict(raw_gen_state_dict, strict=True) + loaded_prequantized_generator = True + prequantized_generator_backend = "fouroversix" + if is_main_process: + print( + f"[NVFP4] Prepared SP generator: {len(dropped_modules)} materialized modules, " + f"{len(matched_modules)} modules excluded" + ) + elif config.use_ema: + missing, unexpected = pipeline.generator.load_state_dict(raw_gen_state_dict, strict=False) + if is_main_process and (missing or unexpected): + print(f"[SP][warn] missing={len(missing)}, unexpected={len(unexpected)}") + else: + pipeline.generator.load_state_dict(raw_gen_state_dict, strict=True) + +pipeline.is_lora_enabled = False +pipeline.is_lora_merged = False +if loaded_prequantized_generator: + has_lora_adapter = False + merge_lora = False + config.merge_lora = False + +if getattr(config, "model_quant", False) and not merge_lora and not loaded_prequantized_generator: + pipeline.generator.model, materialize_quantized_weights_for_inference = quantize_generator_model( + pipeline.generator.model, + config=config, + keep_master_weights=has_lora_adapter, + is_main_process=is_main_process, + ) + +if has_lora_adapter: + if is_main_process: + print(f"[SP] Applying LoRA config: {config.adapter}") + pipeline.generator.model = configure_lora_for_model( + pipeline.generator.model, + model_name="generator", + lora_config=config.adapter, + is_main_process=is_main_process, + ) + lora_ckpt_path = getattr(config, "lora_ckpt", None) + if lora_ckpt_path: + lora_checkpoint = torch.load(lora_ckpt_path, map_location="cpu", mmap=True) + if isinstance(lora_checkpoint, dict) and "generator_lora" in lora_checkpoint: + peft.set_peft_model_state_dict(pipeline.generator.model, lora_checkpoint["generator_lora"]) + else: + peft.set_peft_model_state_dict(pipeline.generator.model, lora_checkpoint) + elif generator_lora_state is not None: + peft.set_peft_model_state_dict(pipeline.generator.model, generator_lora_state) + if merge_lora: + pipeline.generator.model = pipeline.generator.model.merge_and_unload(safe_merge=True) + pipeline.is_lora_merged = True + else: + pipeline.is_lora_enabled = True +elif merge_lora and is_main_process: + print("merge_lora=True requested but no adapter config was found; continuing without LoRA merge") + +del generator_checkpoint + +if loaded_prequantized_generator: + pipeline.text_encoder.to(dtype=torch.bfloat16) + pipeline.vae.to(dtype=torch.bfloat16) +else: + pipeline = pipeline.to(dtype=torch.bfloat16) +if low_memory: + DynamicSwapInstaller.install_model(pipeline.text_encoder, device=device) +pipeline.generator.to(device=device) + +if getattr(config, "model_quant", False) and not loaded_prequantized_generator: + if merge_lora: + pipeline.generator.model, materialize_quantized_weights_for_inference = quantize_generator_model( + pipeline.generator.model, + config=config, + keep_master_weights=False, + is_main_process=is_main_process, + ) + stage_desc = "after LoRA merge" if pipeline.is_lora_merged else "for inference" + else: + stage_desc = "after LoRA wrapping" if pipeline.is_lora_enabled else "for inference" + materialize_quantized_generator( + pipeline.generator.model, + device=device, + materialize_fn=materialize_quantized_weights_for_inference, + stage_desc=stage_desc, + is_main_process=is_main_process, + ) +elif loaded_prequantized_generator and is_main_process: + print(f"[NVFP4] Using pre-saved {prequantized_generator_backend} generator weights") + +pipeline.generator.model.eval().requires_grad_(False) +configure_generator_torch_compile(pipeline, config, is_main_process) + +vae_device_str = getattr(config, "vae_device", None) +use_dedicated_vae_device = bool(getattr(config, "streaming_vae", False)) and bool(vae_device_str) +if use_dedicated_vae_device and sp_rank == 0: + vae_device = torch.device(vae_device_str) + pipeline.vae.to(device="cpu") + pipeline.vae.to(device=vae_device) + if hasattr(pipeline.vae, "mean"): + pipeline.vae.mean = pipeline.vae.mean.to(device=vae_device) + pipeline.vae.std = pipeline.vae.std.to(device=vae_device) + if is_main_process: + print(f"[SP] VAE on {vae_device}, diffusion on {device}") +else: + pipeline.vae.to(device=device) + +nfpb = getattr(config, "num_frame_per_block", 8) +num_blocks = config.num_output_frames // nfpb +dataset = MultiTextConcatDataset( + data_path=config.data_path, + num_blocks=num_blocks, + chunks_per_shot=getattr(config, "chunks_per_shot", 0), + scene_cut_prefix=getattr(config, "scene_cut_prefix", "The scene transitions. "), + deterministic=True, +) +if is_main_process: + print(f"[data] data_path={config.data_path}, mode={dataset._mode}, num_blocks={num_blocks}") +num_prompts = len(dataset) +if use_multi_dp: + sampler = DistributedSampler( + dataset, + num_replicas=total_dp_groups, + rank=dp_rank, + shuffle=False, + drop_last=True, + ) +elif dist.is_initialized(): + sampler = SequentialSampler(dataset) +else: + sampler = SequentialSampler(dataset) +dataloader = DataLoader( + dataset, batch_size=1, sampler=sampler, num_workers=0, + drop_last=False, collate_fn=eval_collate_fn, +) + +if is_main_process: + os.makedirs(config.output_folder, exist_ok=True) +if dist.is_initialized(): + dist.barrier() + +save_latents_only = section_get( + config, + "inference", + "save_latents_only", + getattr(config, "save_latents_only", getattr(config, "save_latent_only", False)), + aliases=("save_latent_only", "return_latents"), +) + +for i, batch_data in tqdm(enumerate(dataloader), disable=not is_main_process): + idx = batch_data["idx"].item() + block_prompts = list(batch_data["prompts"][0]) + if len(block_prompts) < num_blocks: + block_prompts += [block_prompts[-1]] * (num_blocks - len(block_prompts)) + elif len(block_prompts) > num_blocks: + block_prompts = block_prompts[:num_blocks] + prompt = block_prompts[0] + prompts = [block_prompts] * config.num_samples + + shape = config.image_or_video_shape + sampled_noise = torch.randn( + [config.num_samples, config.num_output_frames, shape[2], shape[3], shape[4]], + device=device, + dtype=torch.bfloat16, + ) + if use_effective_sp: + src = dist.get_global_rank(sp_group, 0) + dist.broadcast(sampled_noise, src=src, group=sp_group) + + if is_main_process: + print(f"\n[SP] Generating video {idx}: {prompt[:60]}...") + generated = pipeline.inference( + noise=sampled_noise, + text_prompts=prompts, + return_latents=save_latents_only, + ) + + should_save = (sp_rank == 0) if use_effective_sp else True + if idx < num_prompts and should_save: + if getattr(pipeline, "is_lora_merged", False): + model_type = "merged_lora" + elif getattr(pipeline, "is_lora_enabled", False): + model_type = "lora" + elif getattr(config, "use_ema", False): + model_type = "ema" + else: + model_type = "regular" + mode = f"dp{dp_rank}_sp{effective_sp_size}" if use_multi_dp else f"sp{effective_sp_size}" + if save_latents_only: + latents = generated + else: + current_video = rearrange(generated, "b t c h w -> b t h w c").cpu() + video = 255.0 * current_video + if hasattr(pipeline.vae, "model") and hasattr(pipeline.vae.model, "clear_cache"): + pipeline.vae.model.clear_cache() + + for seed_idx in range(config.num_samples): + if config.save_with_index: + base_name = f"rank{rank}-{idx}-{seed_idx}_{model_type}_{mode}" + else: + base_name = f"rank{rank}-{prompt[:100]}-{seed_idx}_{model_type}_{mode}" + if save_latents_only: + torch.save(latents[seed_idx].cpu(), os.path.join(config.output_folder, f"{base_name}.pt")) + else: + output_path = os.path.join(config.output_folder, f"{base_name}.mp4") + fps = 24 if "5B" in config.model_kwargs.model_name else 16 + write_video(output_path, video[seed_idx], fps=fps) + if is_main_process: + print(f"[SP] Saved: {output_path}") + save_prompts_to_txt( + prompts[seed_idx] if isinstance(prompts[seed_idx], list) else [prompts[seed_idx]], + os.path.join(config.output_folder, f"{base_name}_prompts.txt"), + is_main_process=is_main_process, + ) + + if config.inference_iter != -1 and i >= config.inference_iter: + break + +if dist.is_initialized(): + dist.barrier() + dist.destroy_process_group() diff --git a/model/__init__.py b/model/__init__.py new file mode 100644 index 0000000..1d37685 --- /dev/null +++ b/model/__init__.py @@ -0,0 +1,6 @@ +from .dmd import DMD +from .diffusion import CausalDiffusion +__all__ = [ + "DMD", + "CausalDiffusion", +] diff --git a/model/base.py b/model/base.py new file mode 100644 index 0000000..1ea3eb8 --- /dev/null +++ b/model/base.py @@ -0,0 +1,545 @@ +# Adopted from https://github.com/guandeh17/Self-Forcing +# SPDX-License-Identifier: Apache-2.0 +from typing import Tuple +from einops import rearrange +from torch import nn +import torch.distributed as dist +import torch +import math + +from pipeline import SelfForcingTrainingPipeline +from utils.config import section_get +from utils.loss import get_denoising_loss +from utils.wan_5b_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper + + +def build_default_denoising_step_list(sampling_steps, num_train_timesteps=1000, shift=1.0, include_zero=True): + sigmas = torch.linspace(1.0, 0.0, int(sampling_steps) + 1, dtype=torch.float32)[:-1] + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + timesteps = (sigmas * num_train_timesteps).to(torch.long) + if include_zero: + timesteps = torch.cat([timesteps, torch.zeros(1, dtype=torch.long)]) + return timesteps + + +class BaseModel(nn.Module): + def __init__(self, args, device): + super().__init__() + print("args.model_kwargs.model_name", args.model_kwargs.model_name) + self._initialize_models(args, device) + + self.device = device + self.args = args + self.independent_first_frame = getattr(args, "independent_first_frame", False) + self.dtype = torch.bfloat16 if args.mixed_precision else torch.float32 + self.denoising_step_list = None + if getattr(args, "denoising_step_list", None) is not None: + self.denoising_step_list = torch.tensor(args.denoising_step_list, dtype=torch.long) + if getattr(args, "warp_denoising_step", False): + timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32))) + self.denoising_step_list = timesteps[1000 - self.denoising_step_list] + elif getattr(args, "sampling_steps", None) is not None: + self.denoising_step_list = build_default_denoising_step_list( + sampling_steps=args.sampling_steps, + num_train_timesteps=getattr(args, "num_train_timestep", self.scheduler.num_train_timesteps), + shift=getattr(args, "timestep_shift", self.scheduler.shift), + include_zero=True, + ) + + def _initialize_models(self, args, device): + self.real_model_name = getattr(args, "real_name", "Wan2.2-TI2V-5B") + self.fake_model_name = getattr(args, "fake_name", "Wan2.2-TI2V-5B") + self.local_attn_size = section_get( + args, + "inference", + "local_attn_size", + getattr(args, "model_kwargs", {}).get("local_attn_size", -1), + aliases=("inference_local_attn_size",), + ) + all_causal = getattr(args, "all_causal", False) + score_is_causal = all_causal + + model_name = args.model_kwargs.get("model_name", "Wan2.2-TI2V-5B") + if "5B" not in model_name: + raise ValueError(f"Only Wan2.2-TI2V-5B is supported in this release, got {model_name}") + if not dist.is_initialized() or dist.get_rank() == 0: + tag = "all-causal 5B mode" if all_causal else "Wan2.2-TI2V-5B" + print(f"Using {tag}") + + # Generator + generator_is_causal = getattr(args, "generator_is_causal", True) + self.generator = WanDiffusionWrapper(**getattr(args, "model_kwargs", {}), is_causal=generator_is_causal) + self.generator.model.requires_grad_(True) + + # Real Score + real_kwargs = getattr(args, "real_model_kwargs", {"model_name": self.real_model_name}) + self.real_score = WanDiffusionWrapper(**real_kwargs, is_causal=score_is_causal) + self.real_score.model.requires_grad_(False) + + # Fake Score + fake_kwargs = getattr(args, "fake_model_kwargs", {"model_name": self.fake_model_name}) + self.fake_score = WanDiffusionWrapper(**fake_kwargs, is_causal=score_is_causal) + self.fake_score.model.requires_grad_(True) + + # Text Encoder & VAE + self.text_encoder = WanTextEncoder() + self.text_encoder.requires_grad_(False) + + self.vae = WanVAEWrapper() + self.vae.requires_grad_(False) + + self.scheduler = self.generator.get_scheduler() + self.scheduler.timesteps = self.scheduler.timesteps.to(device) + + def _get_timestep( + self, + min_timestep: int, + max_timestep: int, + batch_size: int, + num_frame: int, + num_frame_per_block: int, + uniform_timestep: bool = False + ) -> torch.Tensor: + """ + Randomly generate a timestep tensor based on the generator's task type. It uniformly samples a timestep + from the range [min_timestep, max_timestep], and returns a tensor of shape [batch_size, num_frame]. + - If uniform_timestep, it will use the same timestep for all frames. + - If not uniform_timestep, it will use a different timestep for each block. + """ + if uniform_timestep: + timestep = torch.randint( + min_timestep, + max_timestep, + [batch_size, 1], + device=self.device, + dtype=torch.long + ).repeat(1, num_frame) + return timestep + else: + timestep = torch.randint( + min_timestep, + max_timestep, + [batch_size, num_frame], + device=self.device, + dtype=torch.long + ) + # make the noise level the same within every block + if self.independent_first_frame and not getattr(self.args, "i2v", False): + # the first frame is always kept the same + timestep_from_second = timestep[:, 1:] + timestep_from_second = timestep_from_second.reshape( + timestep_from_second.shape[0], -1, num_frame_per_block) + timestep_from_second[:, :, 1:] = timestep_from_second[:, :, 0:1] + timestep_from_second = timestep_from_second.reshape( + timestep_from_second.shape[0], -1) + timestep = torch.cat([timestep[:, 0:1], timestep_from_second], dim=1) + else: + timestep = timestep.reshape( + timestep.shape[0], -1, num_frame_per_block) + timestep[:, :, 1:] = timestep[:, :, 0:1] + timestep = timestep.reshape(timestep.shape[0], -1) + return timestep + + +class SelfForcingModel(BaseModel): + def __init__(self, args, device): + super().__init__(args, device) + self.denoising_loss_func = get_denoising_loss(getattr(args, "denoising_loss_type", "flow"))() + + def _run_generator( + self, + image_or_video_shape, + conditional_dict: dict, + initial_latent: torch.tensor = None, + slice_last_frames: int = 21, + noise=None, + clean_latent: torch.Tensor = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Optionally simulate the generator's input from noise using backward simulation + and then run the generator for one-step. + Input: + - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W]. + - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings). + - initial_latent: a tensor containing the initial latents [B, F, C, H, W]. + - slice_last_frames: number of frames to keep from the end. + - noise: optional pre-sampled noise. + - clean_latent: a tensor [B, F, C, H, W] for off-policy mode (backward_simulation=False). + Output: + - pred_image: a tensor with shape [B, F, C, H, W]. + - gradient_mask: boolean tensor or None. + - denoised_timestep_from: int or None. + - denoised_timestep_to: int or None. + """ + use_backward_simulation = getattr(self.args, "backward_simulation", True) + + if use_backward_simulation: + return self._run_generator_backward_simulation( + image_or_video_shape=image_or_video_shape, + conditional_dict=conditional_dict, + initial_latent=initial_latent, + slice_last_frames=slice_last_frames, + noise=noise, + ) + else: + assert clean_latent is not None, "clean_latent is required when backward_simulation=False" + return self._run_generator_off_policy( + image_or_video_shape=image_or_video_shape, + conditional_dict=conditional_dict, + clean_latent=clean_latent, + initial_latent=initial_latent, + ) + + def _run_generator_off_policy( + self, + image_or_video_shape, + conditional_dict: dict, + clean_latent: torch.Tensor, + initial_latent: torch.Tensor = None, + ) -> Tuple[torch.Tensor, torch.Tensor, int, int]: + """ + Off-policy generator: add noise to clean_latent at each timestep in + denoising_step_list, randomly pick one, and run the generator for a + single forward pass. Returns (pred_x0, gradient_mask, + denoised_timestep_from, denoised_timestep_to). + """ + batch_size, num_frame = image_or_video_shape[:2] + denoising_step_list = self.denoising_step_list.to(self.device) + + # Build noisy versions at every timestep in the schedule + simulated_noisy_input = [] + for ts in denoising_step_list: + rand_noise = torch.randn_like(clean_latent) + noisy_timestep = ts * torch.ones( + [batch_size, num_frame], device=self.device, dtype=torch.long) + + if ts.item() != 0: + noisy_image = self.scheduler.add_noise( + clean_latent.flatten(0, 1), + rand_noise.flatten(0, 1), + noisy_timestep.flatten(0, 1), + ).unflatten(0, (batch_size, num_frame)) + else: + noisy_image = clean_latent + simulated_noisy_input.append(noisy_image) + + simulated_noisy_input = torch.stack(simulated_noisy_input, dim=1) # [B, T, F, C, H, W] + + # Randomly sample a step index [B, F], uniform within each block + num_steps = len(denoising_step_list) + generator_is_causal = getattr(self.args, "generator_is_causal", True) + if not generator_is_causal: + # Bidirectional generator: all frames must share the same timestep + index = torch.randint(0, num_steps, [batch_size, 1], + device=self.device, dtype=torch.long).expand(-1, num_frame).contiguous() + else: + index = torch.randint(0, num_steps, [batch_size, num_frame], + device=self.device, dtype=torch.long) + # Make the index the same within every block + if self.independent_first_frame and not getattr(self.args, "i2v", False): + idx_rest = index[:, 1:] + idx_rest = idx_rest.reshape(batch_size, -1, self.num_frame_per_block) + idx_rest[:, :, 1:] = idx_rest[:, :, 0:1] + index = torch.cat([index[:, :1], idx_rest.reshape(batch_size, -1)], dim=1) + else: + index = index.reshape(batch_size, -1, self.num_frame_per_block) + index[:, :, 1:] = index[:, :, 0:1] + index = index.reshape(batch_size, -1) + + # Gather the noisy input corresponding to the sampled index + noisy_input = torch.gather( + simulated_noisy_input, dim=1, + index=index.reshape(batch_size, 1, num_frame, 1, 1, 1).expand( + -1, -1, -1, *image_or_video_shape[2:]) + ).squeeze(1) # [B, F, C, H, W] + + timestep = denoising_step_list[index] # [B, F] + context_frames = int(initial_latent.shape[1]) if initial_latent is not None else 0 + if context_frames > 0: + if context_frames >= num_frame: + raise ValueError( + f"initial_latent has {context_frames} frames but training clip has {num_frame}." + ) + noisy_input[:, :context_frames] = initial_latent.to( + device=noisy_input.device, + dtype=noisy_input.dtype, + ) + timestep[:, :context_frames] = 0 + + # Single forward pass through the generator + _, pred_x0 = self.generator( + noisy_image_or_video=noisy_input, + conditional_dict=conditional_dict, + timestep=timestep.float(), + clean_x=clean_latent if getattr(self.args, "teacher_forcing", False) else None, + ) + pred_x0 = pred_x0.to(self.dtype) + + # Derive denoised_timestep_from / to from the sampled index for ts_schedule + # Use the first batch element's first block index as the representative scalar + rep_idx = index[0, 0].item() + denoised_timestep_from = denoising_step_list[rep_idx].item() + if rep_idx + 1 < num_steps: + denoised_timestep_to = denoising_step_list[rep_idx + 1].item() + else: + denoised_timestep_to = 0 + + gradient_mask = None + if context_frames > 0: + pred_x0[:, :context_frames] = initial_latent.to( + device=pred_x0.device, + dtype=pred_x0.dtype, + ) + gradient_mask = torch.ones_like(pred_x0, dtype=torch.bool) + gradient_mask[:, :context_frames] = False + return pred_x0, gradient_mask, denoised_timestep_from, denoised_timestep_to + + def _run_generator_backward_simulation( + self, + image_or_video_shape, + conditional_dict: dict, + initial_latent: torch.tensor = None, + slice_last_frames: int = 21, + noise=None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + On-policy generator via backward simulation (original path). + """ + if initial_latent is not None: + conditional_dict["initial_latent"] = initial_latent + noise_shape = image_or_video_shape.copy() + + separate_first_frame = self.independent_first_frame and not getattr(self.args, "i2v", False) + min_num_frames = (self.min_num_training_frames - 1) if separate_first_frame else self.min_num_training_frames + max_num_frames = self.num_training_frames - 1 if separate_first_frame else self.num_training_frames + assert max_num_frames % self.num_frame_per_block == 0 + assert min_num_frames % self.num_frame_per_block == 0 + max_num_blocks = max_num_frames // self.num_frame_per_block + min_num_blocks = min_num_frames // self.num_frame_per_block + num_generated_blocks = torch.randint(min_num_blocks, max_num_blocks + 1, (1,), device=self.device) + dist.broadcast(num_generated_blocks, src=0) + num_generated_blocks = num_generated_blocks.item() + num_generated_frames = num_generated_blocks * self.num_frame_per_block + if separate_first_frame and initial_latent is None: + num_generated_frames += 1 + min_num_frames += 1 + noise_shape[1] = num_generated_frames + if noise is not None: + noise = noise[:, :num_generated_frames] + else: + noise = torch.randn(noise_shape, device=self.device, dtype=self.dtype) + + pred_image_or_video, denoised_timestep_from, denoised_timestep_to = self._consistency_backward_simulation( + noise=noise, + slice_last_frames=slice_last_frames, + **conditional_dict, + ) + + if slice_last_frames != -1 and pred_image_or_video.shape[1] > slice_last_frames: + with torch.no_grad(): + if slice_last_frames > 1: + latent_to_decode = pred_image_or_video[:, :-(slice_last_frames - 1), ...] + else: + latent_to_decode = pred_image_or_video + pixels = self.vae.decode_to_pixel(latent_to_decode) + frame = pixels[:, -1:, ...].to(self.dtype) + frame = rearrange(frame, "b t c h w -> b c t h w") + image_latent = self.vae.encode_to_latent(frame).to(self.dtype) + if slice_last_frames > 1: + last_frames = pred_image_or_video[:, -(slice_last_frames - 1):, ...] + pred_image_or_video_sliced = torch.cat([image_latent, last_frames], dim=1) + else: + pred_image_or_video_sliced = image_latent + if num_generated_frames != min_num_frames: + gradient_mask = torch.ones_like(pred_image_or_video_sliced, dtype=torch.bool) + if self.independent_first_frame: + gradient_mask[:, :1] = False + else: + gradient_mask[:, :self.num_frame_per_block] = False + else: + gradient_mask = None + else: + pred_image_or_video_sliced = pred_image_or_video + if num_generated_frames != min_num_frames: + gradient_mask = torch.ones_like(pred_image_or_video_sliced, dtype=torch.bool) + if self.independent_first_frame: + gradient_mask[:, :1] = False + else: + gradient_mask[:, :self.num_frame_per_block] = False + else: + gradient_mask = None + + pred_image_or_video_sliced = pred_image_or_video_sliced.to(self.dtype) + return pred_image_or_video_sliced, gradient_mask, denoised_timestep_from, denoised_timestep_to + + def _consistency_backward_simulation( + self, + noise: torch.Tensor, + slice_last_frames: int = 21, + **conditional_dict: dict + ) -> torch.Tensor: + """ + Simulate the generator's input from noise to avoid training/inference mismatch. + See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details. + Here we use the consistency sampler (https://arxiv.org/abs/2303.01469) + Input: + - noise: a tensor sampled from N(0, 1) with shape [B, F, C, H, W] where the number of frame is 1 for images. + - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings). + Output: + - output: a tensor with shape [B, T, F, C, H, W]. + T is the total number of timesteps. output[0] is a pure noise and output[i] and i>0 + represents the x0 prediction at each timestep. + """ + generator_is_causal = getattr(self.args, "generator_is_causal", True) + if not generator_is_causal: + return self._bidirectional_backward_simulation( + noise=noise, + slice_last_frames=slice_last_frames, + **conditional_dict, + ) + + if self.inference_pipeline is None: + self._initialize_inference_pipeline() + + return self.inference_pipeline.inference_with_trajectory( + noise=noise, **conditional_dict, slice_last_frames=slice_last_frames + ) + + def _bidirectional_backward_simulation( + self, + noise: torch.Tensor, + slice_last_frames: int = 21, + **conditional_dict: dict + ) -> Tuple[torch.Tensor, int, int]: + """ + Backward simulation for bidirectional (non-causal) generator. + All frames are processed at once at each denoising step — no KV cache, + no block-by-block processing. + """ + from wan_5b.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler + + batch_size, num_frames = noise.shape[:2] + + # Resolve single-segment prompt for bidirectional model + prompt_embeds = conditional_dict["prompt_embeds"] + num_segments = prompt_embeds.shape[0] // batch_size + if num_segments > 1: + prompt_embeds = prompt_embeds.reshape( + batch_size, num_segments, *prompt_embeds.shape[1:])[:, 0] + cond = {**conditional_dict, "prompt_embeds": prompt_embeds} + + # Setup UniPC scheduler + sampling_steps = getattr(self.args, "sampling_steps", None) or len(self.denoising_step_list) + shift = self.scheduler.shift + num_train_timesteps = self.scheduler.num_train_timesteps + sample_scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=num_train_timesteps, shift=1, use_dynamic_shifting=False) + sample_scheduler.set_timesteps(sampling_steps, device=noise.device, shift=shift) + unipc_timesteps = sample_scheduler.timesteps + num_denoising_steps = len(unipc_timesteps) + + # Pick a random exit step (synchronized across ranks) + last_step_only = getattr(self.args, "last_step_only", False) + if last_step_only: + exit_step = num_denoising_steps - 1 + else: + exit_step_t = torch.randint(0, num_denoising_steps, (1,), device=self.device) + dist.broadcast(exit_step_t, src=0) + exit_step = exit_step_t.item() + + # Multi-step denoising loop (full-sequence bidirectional forward each step) + latents = noise + for index, t in enumerate(unipc_timesteps): + timestep = t * torch.ones( + [batch_size, num_frames], device=noise.device, dtype=torch.float32) + + if index < exit_step: + with torch.no_grad(): + flow_pred, _ = self.generator( + noisy_image_or_video=latents, + conditional_dict=cond, + timestep=timestep, + ) + latents = sample_scheduler.step( + flow_pred, t, latents, return_dict=False)[0] + else: + # Exit step: forward with gradient + flow_pred, denoised_pred = self.generator( + noisy_image_or_video=latents, + conditional_dict=cond, + timestep=timestep, + ) + break + + # Compute denoised_timestep_from / to + if exit_step == num_denoising_steps - 1: + denoised_timestep_to = 0 + denoised_timestep_from = 1000 - torch.argmin( + (self.scheduler.timesteps.cuda() - unipc_timesteps[exit_step].cuda()).abs(), dim=0).item() + else: + denoised_timestep_to = 1000 - torch.argmin( + (self.scheduler.timesteps.cuda() - unipc_timesteps[exit_step + 1].cuda()).abs(), dim=0).item() + denoised_timestep_from = 1000 - torch.argmin( + (self.scheduler.timesteps.cuda() - unipc_timesteps[exit_step].cuda()).abs(), dim=0).item() + + return denoised_pred, denoised_timestep_from, denoised_timestep_to + + def _initialize_inference_pipeline(self): + """ + Lazy initialize the inference pipeline during the first backward simulation run. + Here we encapsulate the inference code with a model-dependent outside function. + We pass our FSDP-wrapped modules into the pipeline to save memory. + """ + local_attn_size = section_get( + self.args, + "inference", + "local_attn_size", + getattr(self.args, "model_kwargs", {}).get("local_attn_size", -1), + aliases=("inference_local_attn_size",), + ) + sink_size = section_get( + self.args, + "inference", + "sink_size", + getattr(self.args, "model_kwargs", {}).get("sink_size", 0), + aliases=("inference_sink_size",), + ) + multi_shot_sink = section_get(self.args, "inference", "multi_shot_sink", False) + multi_shot_rope_offset = section_get( + self.args, + "inference", + "multi_shot_rope_offset", + 0.0, + ) + scene_cut_prefix = section_get(self.args, "inference", "scene_cut_prefix", "[SCENE_CUT]") + slice_last_frames = getattr(self.args, "slice_last_frames", 21) + # do not use self.num_training_frames, because it is changed by generator_loss and critic_loss + num_training_frames = getattr(self.args, "num_training_frames") + if local_attn_size == -1: + kv_cache_size = num_training_frames + else: + kv_cache_size = min(local_attn_size + slice_last_frames, num_training_frames) + frame_seq_length = math.prod(self.args.image_or_video_shape[-2:]) // 4 + self.inference_pipeline = SelfForcingTrainingPipeline( + denoising_step_list=self.denoising_step_list, + scheduler=self.scheduler, + generator=self.generator, + num_frame_per_block=self.num_frame_per_block, + independent_first_frame=self.independent_first_frame, + same_step_across_blocks=getattr( + self.args, "same_step_across_blocks", getattr(self, "same_step_across_blocks", False) + ), + last_step_only=getattr(self.args, "last_step_only", False), + num_max_frames=kv_cache_size, + context_noise=getattr(self.args, "context_noise", 0), + sampling_steps=getattr(self.args, "sampling_steps", None), + local_attn_size=local_attn_size, + sink_size=sink_size, + multi_shot_sink=multi_shot_sink, + scene_cut_prefix=scene_cut_prefix, + multi_shot_rope_offset=multi_shot_rope_offset, + slice_last_frames=slice_last_frames, + num_training_frames=num_training_frames, + model_name=getattr(self.args, "model_kwargs", {}).get("model_name", None), + frame_seq_length=frame_seq_length, + ) diff --git a/model/diffusion.py b/model/diffusion.py new file mode 100644 index 0000000..c00b0a8 --- /dev/null +++ b/model/diffusion.py @@ -0,0 +1,574 @@ +# Adopted from https://github.com/guandeh17/Self-Forcing +# SPDX-License-Identifier: Apache-2.0 + +from typing import Tuple +import random +import torch + +from model.base import BaseModel +from pipeline import CausalDiffusionInferencePipeline +from utils.i2v_conditioning import _overwrite_i2v_context, _zero_i2v_context_timestep +from utils.wan_5b_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper + + +class CausalDiffusion(BaseModel): + def __init__(self, args, device): + """ + Initialize the Diffusion loss module. + """ + super().__init__(args, device) + self.num_frame_per_block = getattr(args, "num_frame_per_block", 1) + if self.num_frame_per_block > 1: + self.generator.model.num_frame_per_block = self.num_frame_per_block + self.independent_first_frame = getattr(args, "independent_first_frame", False) + if self.independent_first_frame and not getattr(args, "i2v", False): + self.generator.model.independent_first_frame = True + + if args.gradient_checkpointing: + self.generator.enable_gradient_checkpointing() + + # Step 2: Initialize all hyperparameters + self.num_train_timestep = args.num_train_timestep + self.min_step = int(0.02 * self.num_train_timestep) + self.max_step = int(0.98 * self.num_train_timestep) + self.guidance_scale = args.guidance_scale + self.timestep_shift = getattr(args, "timestep_shift", 1.0) + self.teacher_forcing = getattr(args, "teacher_forcing", False) + # Noise augmentation in teacher forcing, we add small noise to clean context latents + self.noise_augmentation_max_timestep = getattr(args, "noise_augmentation_max_timestep", 0) + + self.args = args + self.device = device + self.inference_pipeline = None + + # Error recycling (SVI-style error buffer) + # When ``enable_position_bucketing`` is true, each rank holds a 2D + # buffer ``(local_block_position × timestep)``. The pos dimension only + # covers the LOCAL slice of the sequence this rank is responsible for + # (no cross-SP-rank pos sharing — those positions are simply not + # reachable by this rank during forward), so memory cost scales as + # ``num_blocks_global / sp_size`` instead of ``num_blocks_global``. + # ``global_block_offset`` is recorded for logging only. + # During the first ``buffer_warmup_iter`` global steps, errors are + # all-gathered across the DP group (ranks with the same SP rank but + # different DP replicas), so each rank's local pos buckets fill up + # ``dp_size`` × faster without any wasted bandwidth. + self.error_buffer = None + self.noise_error_buffer = None + self.er_num_blocks = 0 # local; >0 means 2D position-bucketed + self.er_block_offset = 0 # global block offset for THIS rank + er_cfg = getattr(args, "error_recycling", None) + if er_cfg is not None and getattr(er_cfg, "enabled", False): + from utils.error_buffer import build_error_buffer + cfg_dict = er_cfg if isinstance(er_cfg, dict) else dict(er_cfg) + cfg_dict.setdefault("num_train_timesteps", self.num_train_timestep) + sp_size = int(getattr(args, "sequence_parallel_size", 1) or 1) + if cfg_dict.get("enable_position_bucketing", False): + shape = list(getattr(args, "image_or_video_shape", [1, 0])) + total_frames = int(shape[1]) if len(shape) > 1 else 0 + assert total_frames > 0 and self.num_frame_per_block > 0, ( + "enable_position_bucketing=true requires " + "image_or_video_shape[1] and num_frame_per_block to be set." + ) + num_blocks_global = total_frames // self.num_frame_per_block + assert num_blocks_global % sp_size == 0, ( + f"num_blocks_global ({num_blocks_global}) must be divisible " + f"by sequence_parallel_size ({sp_size})." + ) + self.er_num_blocks = num_blocks_global // sp_size # local + # Determine this rank's SP index → global block offset. + if sp_size > 1: + import torch.distributed as dist + if dist.is_initialized(): + sp_rank = dist.get_rank() % sp_size + else: + sp_rank = 0 + else: + sp_rank = 0 + self.er_block_offset = sp_rank * self.er_num_blocks + + # Shard timestep buckets across SP ranks: each SP rank only + # stores t_bucket % sp_size == sp_rank, cutting per-rank CPU + # memory by ~sp_size. This uses the same SP dimension that + # already splits positions in 2D mode, so both 1D and 2D + # follow one save/load pattern (per sp_rank). + import torch.distributed as dist + if sp_size > 1 and dist.is_initialized(): + er_shard_rank = dist.get_rank() % sp_size + er_shard_size = sp_size + else: + er_shard_rank = 0 + er_shard_size = 1 + + self.error_buffer = build_error_buffer( + cfg_dict, num_blocks=self.er_num_blocks, + global_block_offset=self.er_block_offset, + shard_rank=er_shard_rank, shard_size=er_shard_size, + ) + self.noise_error_buffer = build_error_buffer( + cfg_dict, num_blocks=self.er_num_blocks, + global_block_offset=self.er_block_offset, + shard_rank=er_shard_rank, shard_size=er_shard_size, + ) + self.er_context_inject_prob = float(cfg_dict.get("context_inject_prob", 0.9)) + self.er_latent_inject_prob = float(cfg_dict.get("latent_inject_prob", 0.0)) + self.er_noise_inject_prob = float(cfg_dict.get("noise_inject_prob", 0.0)) + self.er_clean_prob = float(cfg_dict.get("clean_prob", 0.0)) + self.er_clean_buffer_update_prob = float(cfg_dict.get("clean_buffer_update_prob", 0.1)) + self.er_start_step = int(cfg_dict.get("start_step", 0)) + self.er_buffer_warmup_iter = int(cfg_dict.get("buffer_warmup_iter", 50)) + self.er_skip_block_0 = bool(cfg_dict.get("skip_block_0", False)) + + def _initialize_models(self, args, device): + model_name = getattr(args.model_kwargs, "model_name", "Wan2.2-TI2V-5B") + if "5B" not in model_name: + raise ValueError(f"Only Wan2.2-TI2V-5B is supported in this release, got {model_name}") + self.generator = WanDiffusionWrapper(**getattr(args, "model_kwargs", {}), is_causal=True) + self.generator.model.requires_grad_(True) + + self.text_encoder = WanTextEncoder() + self.text_encoder.requires_grad_(False) + + self.vae = WanVAEWrapper() + self.vae.requires_grad_(False) + + self.scheduler = self.generator.get_scheduler() + self.scheduler.timesteps = self.scheduler.timesteps.to(device) + + def generator_loss( + self, + image_or_video_shape, + conditional_dict: dict, + unconditional_dict: dict, + clean_latent: torch.Tensor, + initial_latent: torch.Tensor = None, + loss_mask: torch.Tensor = None, + loss_mask_global_valid_count: torch.Tensor = None, + global_step: int = None, + ) -> Tuple[torch.Tensor, dict]: + """ + Generate image/videos from noise and compute the DMD loss. + The noisy input to the generator is backward simulated. + This removes the need of any datasets during distillation. + See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details. + Input: + - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W]. + - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings). + - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings). + - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used. + - loss_mask: optional tensor of shape [B, F] with 1.0 for valid frames and 0.0 for padded frames. + Under Sequence Parallel this is already the local chunk. + - loss_mask_global_valid_count: optional scalar tensor with the total valid count across all SP ranks. + When provided (SP mode), used as the denominator instead of the local loss_mask.sum(). + - global_step: current training step, used for error recycling delayed start. + Output: + - loss: a scalar tensor representing the generator loss. + - generator_log_dict: a dictionary containing the intermediate tensors for logging. + """ + batch_size, num_frame = image_or_video_shape[:2] + + noise = torch.randn_like(clean_latent) + # Step 2: Randomly sample a timestep and add noise to denoiser inputs + index = self._get_timestep( + 0, + self.scheduler.num_train_timesteps, + image_or_video_shape[0], + image_or_video_shape[1], + self.num_frame_per_block, + uniform_timestep=False + ) + timestep = self.scheduler.timesteps[index].to(dtype=self.dtype, device=self.device) + context_latent = ( + initial_latent + if getattr(self.args, "i2v", False) and initial_latent is not None + else None + ) + context_frames = int(context_latent.shape[1]) if context_latent is not None else 0 + if context_frames > 0: + if context_frames >= num_frame: + raise ValueError( + f"initial_latent has {context_frames} frames but training clip has {num_frame}." + ) + timestep[:, :context_frames] = 0 + + # Step 2.5 & 3.5: Error recycling — clean_prob acts as a master switch. + # When clean_prob fires, skip ALL error injection and use pristine input; + # otherwise each injection type rolls its own probability and is gated + # only by whether the corresponding buffer has any samples (SVI behavior). + # + # NOTE on rank-sync: random.random() below is INTENTIONALLY independent + # across ranks. None of these decisions guard a collective call (we use + # SVI's pattern of unconditional all_gather + local random replay in + # Step 5), so per-rank divergence here only affects which slice of data + # gets corrupted on which rank — perfectly safe under DP+SP, and matches + # SVI's behavior exactly. + er_latent_injected = False + er_noise_injected = False + er_injected = False + er_use_clean = False + er_ready = ( + self.error_buffer is not None + and (global_step is None or global_step >= self.er_start_step) + ) + if er_ready and self.er_clean_prob > 0 and random.random() < self.er_clean_prob: + er_use_clean = True + + # Noise error injection (SVI's noise_prob): corrupt noise input. + # training_target is then computed with the corrupted noise so the + # model learns to predict a self-correcting velocity (SVI Eq. logic). + noise_for_train = noise + if ( + er_ready and not er_use_clean + and self.er_noise_inject_prob > 0 + and not self.noise_error_buffer.is_empty() + and random.random() < self.er_noise_inject_prob + ): + noise_for_train = self._inject_noise_error_buffer( + noise, index, batch_size, num_frame + ) + er_noise_injected = True + + # Latent error injection (SVI's latent_prob): corrupt clean_latent + # before noising. training_target keeps pointing to ORIGINAL clean_latent. + clean_latent_for_noise = clean_latent + if ( + er_ready and not er_use_clean + and self.er_latent_inject_prob > 0 + and not self.error_buffer.is_empty() + and random.random() < self.er_latent_inject_prob + ): + clean_latent_for_noise = self._inject_latent_error_buffer( + clean_latent, index, batch_size, num_frame + ) + er_latent_injected = True + + noisy_latents = self.scheduler.add_noise( + clean_latent_for_noise.flatten(0, 1), + noise_for_train.flatten(0, 1), + timestep.flatten(0, 1) + ).unflatten(0, (batch_size, num_frame)) + training_target = self.scheduler.training_target(clean_latent, noise_for_train, timestep) + if context_frames > 0: + noisy_latents[:, :context_frames] = context_latent.to( + device=noisy_latents.device, + dtype=noisy_latents.dtype, + ) + training_target[:, :context_frames] = 0 + + # Step 3: Noise augmentation, also add small noise to clean context latents + if self.noise_augmentation_max_timestep > 0: + index_clean_aug = self._get_timestep( + 0, + self.noise_augmentation_max_timestep, + image_or_video_shape[0], + image_or_video_shape[1], + self.num_frame_per_block, + uniform_timestep=False + ) + timestep_clean_aug = self.scheduler.timesteps[index_clean_aug].to(dtype=self.dtype, device=self.device) + clean_latent_aug = self.scheduler.add_noise( + clean_latent.flatten(0, 1), + noise.flatten(0, 1), + timestep_clean_aug.flatten(0, 1) + ).unflatten(0, (batch_size, num_frame)) + else: + clean_latent_aug = clean_latent + timestep_clean_aug = None + + # Step 3.5: Error recycling — inject sampled errors into clean prefix. + # 2D mode: per-position (random timestep). 1D mode: SVI global sampling. + if ( + er_ready and not er_use_clean + and not self.error_buffer.is_empty() + and random.random() < self.er_context_inject_prob + ): + clean_latent_aug = self._inject_error_buffer( + clean_latent_aug, index, batch_size, num_frame + ) + er_injected = True + + if context_frames > 0: + clean_latent_aug = _overwrite_i2v_context( + clean_latent_aug, context_latent, context_frames + ) + if timestep_clean_aug is not None: + timestep_clean_aug = _zero_i2v_context_timestep( + timestep_clean_aug, context_frames + ) + + # Compute loss + flow_pred, x0_pred = self.generator( + noisy_image_or_video=noisy_latents, + conditional_dict=conditional_dict, + timestep=timestep, + clean_x=clean_latent_aug if self.teacher_forcing else None, + aug_t=timestep_clean_aug if self.teacher_forcing else None, + ) + loss = torch.nn.functional.mse_loss( + flow_pred.float(), training_target.float(), reduction='none' + ).mean(dim=(2, 3, 4)) + loss = loss * self.scheduler.training_weight(timestep).unflatten(0, (batch_size, num_frame)) + if context_frames > 0: + if loss_mask is None: + loss_mask = torch.ones( + (batch_size, num_frame), + device=loss.device, + dtype=loss.dtype, + ) + loss_mask[:, :context_frames] = 0 + if loss_mask is not None: + loss = loss * loss_mask + valid_count = loss_mask_global_valid_count if loss_mask_global_valid_count is not None else loss_mask.sum() + loss = loss.sum() / valid_count.clamp(min=1.0) + else: + loss = loss.mean() + + log_dict = { + "x0": clean_latent.detach(), + "x0_pred": x0_pred.detach() + } + + # Step 5: Store prediction errors into error buffer. + # + # SVI-style two-phase pattern (avoids the rank-divergent collective + # deadlock that an ``if random.random() < p: all_gather()`` would + # introduce): + # PHASE A (collective, UNCONDITIONAL): every rank reaches the + # all_gather call together so NCCL stays in sync. + # PHASE B (local, GATED): each rank independently decides whether + # to actually replay the gathered items into its buffer. The + # gate uses random.random() per rank — divergence here is fine + # because no further collective follows. + if self.error_buffer is not None: + with torch.no_grad(): + # latent error: x0_pred - clean_latent ≡ -σ(v_pred - v_gt) — used for context/latent injection + latent_err = x0_pred.detach() - clean_latent.detach() + # noise error: SVI definition is (1-σ)(v_pred - v_gt) so the buffer entry, + # when later added directly to noise, equals ε_pred - ε_gt. + sigma = self.scheduler.sigmas.to(flow_pred.device)[index].reshape( + batch_size, num_frame, 1, 1, 1 + ).to(flow_pred.dtype) + noise_err = (flow_pred.detach() - training_target.detach()) * (1.0 - sigma) + + use_distributed = ( + global_step is not None + and global_step <= self.er_buffer_warmup_iter + ) + + # === PHASE A: collective — runs on EVERY rank, no gating === + if use_distributed: + lat_items = self._gather_errors_for_buffer( + self.error_buffer, latent_err, index, batch_size, num_frame + ) + noise_items = self._gather_errors_for_buffer( + self.noise_error_buffer, noise_err, index, batch_size, num_frame + ) + else: + lat_items = self._collect_local_items( + self.error_buffer, latent_err, index, batch_size, num_frame + ) + noise_items = self._collect_local_items( + self.noise_error_buffer, noise_err, index, batch_size, num_frame + ) + + # === PHASE B: local replay — random.random() per-rank is OK === + # When the input was clean (low-error), only update buffer with + # small probability to avoid flooding it with near-zero samples + # (SVI: clean_buffer_update_prob). + should_update = True + if er_use_clean and random.random() >= self.er_clean_buffer_update_prob: + should_update = False + if should_update: + self._apply_gathered_items(self.error_buffer, lat_items) + self._apply_gathered_items(self.noise_error_buffer, noise_items) + buf_stats = self.error_buffer.stats() + noise_buf_stats = self.noise_error_buffer.stats() + log_dict["er_total_added"] = buf_stats["total_added"] + log_dict["er_filled_buckets"] = buf_stats["filled_buckets"] + log_dict["er_total_entries"] = buf_stats["total_entries"] + log_dict["er_noise_total_entries"] = noise_buf_stats["total_entries"] + log_dict["er_injected"] = er_injected + log_dict["er_latent_injected"] = er_latent_injected + log_dict["er_noise_injected"] = er_noise_injected + + return loss, log_dict + + + def _inject_error_buffer(self, clean_latent_aug, index, batch_size, num_frame): + """Inject errors into the clean prefix (E_img). + + 2D (position-bucketed): the i-th LOCAL prefix block draws from + ``buckets[(i, *)]`` with a RANDOM timestep — the clean prefix is + the product of full ODE integration so its accumulated error can + come from any noise level, but its magnitude scales with the + block's global position. Note ``skip_block_0`` is interpreted in + the GLOBAL frame: only the very first SP rank may skip its block 0. + + 1D (timestep-bucketed): falls back to SVI ``sample_global``. + """ + block_size = self.num_frame_per_block + num_blocks = num_frame // block_size + result = clean_latent_aug.clone() + for b in range(batch_size): + for blk in range(num_blocks): + if self.er_skip_block_0 and (self.er_block_offset + blk) == 0: + continue + if self.er_num_blocks > 0: + err = self.error_buffer.sample_pos_any_t( + blk, device=result.device, dtype=result.dtype + ) + else: + err = self.error_buffer.sample_global( + device=result.device, dtype=result.dtype + ) + if err is not None: + start = blk * block_size + end = start + block_size + result[b, start:end] = result[b, start:end] + err + return result + + def _inject_latent_error_buffer(self, clean_latent, index, batch_size, num_frame): + """Inject errors into clean_latent before noising (E_vid). + + Matches BOTH block_position (LOCAL) and timestep when the buffer is + 2D, else only timestep (SVI default). + """ + block_size = self.num_frame_per_block + num_blocks = num_frame // block_size + index_per_block = index[:, ::block_size] + result = clean_latent.clone() + for b in range(batch_size): + for blk in range(num_blocks): + t_idx = index_per_block[b, blk].item() + pos = blk if self.er_num_blocks > 0 else None + err = self.error_buffer.sample( + t_idx, device=result.device, dtype=result.dtype, + block_pos=pos, + ) + if err is not None: + start = blk * block_size + end = start + block_size + result[b, start:end] = result[b, start:end] + err + return result + + def _inject_noise_error_buffer(self, noise, index, batch_size, num_frame): + """Inject errors into the noise (E_noise). + + Same matching strategy as ``_inject_latent_error_buffer`` but reads + from the dedicated noise buffer. + """ + block_size = self.num_frame_per_block + num_blocks = num_frame // block_size + index_per_block = index[:, ::block_size] + result = noise.clone() + for b in range(batch_size): + for blk in range(num_blocks): + t_idx = index_per_block[b, blk].item() + pos = blk if self.er_num_blocks > 0 else None + err = self.noise_error_buffer.sample( + t_idx, device=result.device, dtype=result.dtype, + block_pos=pos, + ) + if err is not None: + start = blk * block_size + end = start + block_size + result[b, start:end] = result[b, start:end] + err + return result + + def _gather_errors_for_buffer( + self, buffer, error, index, batch_size, num_frame + ): + """All-gather errors/timesteps across the appropriate group and return + a list of ready-to-add ``(err_block, t_idx, pos_or_None)`` items. + + ★ This is a COLLECTIVE — every rank MUST reach this call together. + The caller is responsible for invoking it unconditionally during the + warmup window (just like SVI's ``all_gather`` outside the random + ``if`` blocks). Random decisions about whether to actually consume + the returned items belong to ``_apply_gathered_items`` instead. + + Group selection mirrors SVI's intent: + * **2D (num_blocks > 0)** — DP group only. Other SP ranks' samples + map to positions unreachable by this rank, so cross-SP gather + wastes bandwidth. + * **1D (num_blocks == 0)** — WORLD group (SVI default). Buckets + are pos-agnostic so every rank's errors are valid samples. + """ + import torch.distributed as dist + if not dist.is_initialized() or dist.get_world_size() <= 1: + return self._collect_local_items(buffer, error, index, batch_size, num_frame) + + if buffer.num_blocks > 0: + from wan_5b.distributed.sp_training import get_data_parallel_group + comm_group = get_data_parallel_group() + if comm_group is None: + return self._collect_local_items(buffer, error, index, batch_size, num_frame) + comm_size = dist.get_world_size(comm_group) + else: + comm_group = None + comm_size = dist.get_world_size() + + if comm_size <= 1: + return self._collect_local_items(buffer, error, index, batch_size, num_frame) + + err_local = error.detach().contiguous() + idx_local = index.detach().contiguous() + err_list = [torch.empty_like(err_local) for _ in range(comm_size)] + idx_list = [torch.empty_like(idx_local) for _ in range(comm_size)] + if comm_group is None: + dist.all_gather(err_list, err_local) + dist.all_gather(idx_list, idx_local) + else: + dist.all_gather(err_list, err_local, group=comm_group) + dist.all_gather(idx_list, idx_local, group=comm_group) + + block_size = self.num_frame_per_block + num_blocks = num_frame // block_size + items = [] + for err_r, idx_r in zip(err_list, idx_list): + idx_per_block = idx_r[:, ::block_size] + err_blocks = err_r.reshape( + batch_size, num_blocks, block_size, *err_r.shape[2:] + ) + for b in range(batch_size): + for blk in range(num_blocks): + pos = blk if buffer.num_blocks > 0 else None + items.append((err_blocks[b, blk], idx_per_block[b, blk].item(), pos)) + return items + + def _collect_local_items(self, buffer, error, index, batch_size, num_frame): + """Same item-list format as ``_gather_errors_for_buffer`` but with no + collective — used outside the warmup window or when distributed is off.""" + block_size = self.num_frame_per_block + num_blocks = num_frame // block_size + idx_per_block = index[:, ::block_size] + error_blocks = error.reshape( + batch_size, num_blocks, block_size, *error.shape[2:] + ) + items = [] + for b in range(batch_size): + for blk in range(num_blocks): + pos = blk if buffer.num_blocks > 0 else None + items.append((error_blocks[b, blk], idx_per_block[b, blk].item(), pos)) + return items + + def _apply_gathered_items(self, buffer, items): + """Pure local: drop ``items`` into ``buffer``. No collective, no + cross-rank coordination — each rank may invoke this independently + (or skip it entirely) without risking a deadlock.""" + for err_block, t_idx, pos in items: + buffer.add(err_block, t_idx, block_pos=pos) + + def _initialize_inference_pipeline(self): + """ + Lazy initialize the inference pipeline during the first backward simulation run. + Here we encapsulate the inference code with a model-dependent outside function. + We pass our FSDP-wrapped modules into the pipeline to save memory. + """ + self.inference_pipeline = CausalDiffusionInferencePipeline( + args=self.args, + device=self.device, + generator=self.generator, + text_encoder=self.text_encoder, + vae=self.vae + ) diff --git a/model/dmd.py b/model/dmd.py new file mode 100644 index 0000000..1e59d8b --- /dev/null +++ b/model/dmd.py @@ -0,0 +1,494 @@ +# Adopted from https://github.com/guandeh17/Self-Forcing +# SPDX-License-Identifier: Apache-2.0 + +import torch.nn.functional as F +from typing import Optional, Tuple +import torch +import time + +from model.base import SelfForcingModel +import torch.distributed as dist +from utils.i2v_conditioning import ( + _get_i2v_context_frames, + _i2v_loss_mask_like, + _overwrite_i2v_context, + _zero_i2v_context_timestep, +) + + +class DMD(SelfForcingModel): + def __init__(self, args, device): + """ + Initialize the DMD (Distribution Matching Distillation) module. + This class is self-contained and compute generator and fake score losses + in the forward pass. + """ + super().__init__(args, device) + self.num_frame_per_block = getattr(args, "num_frame_per_block", 1) + self.same_step_across_blocks = getattr(args, "same_step_across_blocks", True) + self.min_num_training_frames = getattr(args, "min_num_training_frames", 21) + self.num_training_frames = getattr(args, "num_training_frames", 21) + + if self.num_frame_per_block > 1: + for diffusion_model in (self.generator, self.real_score, self.fake_score): + if hasattr(diffusion_model.model, "num_frame_per_block"): + diffusion_model.model.num_frame_per_block = self.num_frame_per_block + + self.independent_first_frame = getattr(args, "independent_first_frame", False) + if self.independent_first_frame and not getattr(args, "i2v", False): + self.generator.model.independent_first_frame = True + if args.gradient_checkpointing: + self.generator.enable_gradient_checkpointing() + self.fake_score.enable_gradient_checkpointing() + + # this will be init later with fsdp-wrapped modules + self.inference_pipeline: SelfForcingTrainingPipeline = None + + # Step 2: Initialize all dmd hyperparameters + self.num_train_timestep = getattr(args, "num_train_timestep", 1000) + self.min_step = int(0.02 * self.num_train_timestep) + self.max_step = int(0.98 * self.num_train_timestep) + if hasattr(args, "real_guidance_scale"): + self.real_guidance_scale = args.real_guidance_scale + self.fake_guidance_scale = args.fake_guidance_scale + else: + self.real_guidance_scale = args.guidance_scale + self.fake_guidance_scale = 0.0 + self.timestep_shift = getattr(args, "timestep_shift", 1.0) + self.ts_schedule = getattr(args, "ts_schedule", True) + self.ts_schedule_max = getattr(args, "ts_schedule_max", False) + self.min_score_timestep = getattr(args, "min_score_timestep", 0) + + if getattr(self.scheduler, "alphas_cumprod", None) is not None: + self.scheduler.alphas_cumprod = self.scheduler.alphas_cumprod.to(device) + else: + self.scheduler.alphas_cumprod = None + + @staticmethod + def _slice_block_cond_dict(cond_dict, batch_size, new_num_segments): + """Slice a block-wise conditional dict to keep only the last `new_num_segments` segments.""" + pe = cond_dict["prompt_embeds"] + orig_segs = pe.shape[0] // batch_size + if orig_segs > new_num_segments: + pe = pe.reshape(batch_size, orig_segs, *pe.shape[1:])[:, -new_num_segments:] + return {**cond_dict, "prompt_embeds": pe.reshape(batch_size * new_num_segments, *pe.shape[2:])} + return cond_dict + + def _compute_kl_grad( + self, noisy_image_or_video: torch.Tensor, + estimated_clean_image_or_video: torch.Tensor, + timestep: torch.Tensor, + conditional_dict: dict, unconditional_dict: dict, + normalization: bool = True, + clean_x: Optional[torch.Tensor] = None + ) -> Tuple[torch.Tensor, dict]: + """ + Compute the KL grad (eq 7 in https://arxiv.org/abs/2311.18828). + Input: + - noisy_image_or_video: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images. + - estimated_clean_image_or_video: a tensor with shape [B, F, C, H, W] representing the estimated clean image or video. + - timestep: a tensor with shape [B, F] containing the randomly generated timestep. + - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings). + - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings). + - normalization: a boolean indicating whether to normalize the gradient. + Output: + - kl_grad: a tensor representing the KL grad. + - kl_log_dict: a dictionary containing the intermediate tensors for logging. + """ + # Step 1: Compute the fake score + _, pred_fake_image_cond = self.fake_score( + noisy_image_or_video=noisy_image_or_video, + conditional_dict=conditional_dict, + timestep=timestep, + clean_x=clean_x + ) + + if self.fake_guidance_scale != 0.0: + _, pred_fake_image_uncond = self.fake_score( + noisy_image_or_video=noisy_image_or_video, + conditional_dict=unconditional_dict, + timestep=timestep, + clean_x=clean_x + ) + pred_fake_image = pred_fake_image_cond + ( + pred_fake_image_cond - pred_fake_image_uncond + ) * self.fake_guidance_scale + else: + pred_fake_image = pred_fake_image_cond + + # Step 2: Compute the real score + _, pred_real_image_cond = self.real_score( + noisy_image_or_video=noisy_image_or_video, + conditional_dict=conditional_dict, + timestep=timestep, + clean_x=clean_x + ) + + _, pred_real_image_uncond = self.real_score( + noisy_image_or_video=noisy_image_or_video, + conditional_dict=unconditional_dict, + timestep=timestep, + clean_x=clean_x + ) + + pred_real_image = pred_real_image_cond + ( + pred_real_image_cond - pred_real_image_uncond + ) * self.real_guidance_scale + + # Step 3: Compute the DMD gradient (DMD paper eq. 7). + grad = (pred_fake_image - pred_real_image) + + # NOTE: Changed the normalizer for causal teacher — per-block normalization + if normalization: + p_real = (estimated_clean_image_or_video - pred_real_image) + + B, F, C, H, W = p_real.shape + if dist.get_rank() == 0: + print(f"p_real: {p_real.shape}") + if ( + self.independent_first_frame + and not getattr(self.args, "i2v", False) + and (F - 1) % self.num_frame_per_block == 0 + ): + p_real_tail = p_real[:, 1:] + p_real_blocks = p_real_tail.view( + B, + (F - 1) // self.num_frame_per_block, + self.num_frame_per_block, + C, + H, + W, + ) + normalizer_tail = torch.abs(p_real_blocks).mean(dim=[2, 3, 4, 5], keepdim=True) + normalizer = torch.ones_like(p_real) + normalizer[:, 1:] = normalizer_tail.expand_as(p_real_blocks).reshape( + B, F - 1, C, H, W + ) + else: + p_real_blocks = p_real.view(B, F // self.num_frame_per_block, self.num_frame_per_block, C, H, W) + normalizer = torch.abs(p_real_blocks).mean(dim=[2, 3, 4, 5], keepdim=True) + normalizer = normalizer.expand_as(p_real_blocks).reshape(B, F, C, H, W) + + + grad = grad / normalizer + grad = torch.nan_to_num(grad) + + return grad, { + "dmdtrain_gradient_norm": torch.mean(torch.abs(grad)).detach(), + "timestep": timestep.detach() + } + + def compute_distribution_matching_loss( + self, + image_or_video: torch.Tensor, + conditional_dict: dict, + unconditional_dict: dict, + gradient_mask: Optional[torch.Tensor] = None, + denoised_timestep_from: int = 0, + denoised_timestep_to: int = 0, + clean_x: Optional[torch.Tensor] = None, + initial_latent: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, dict]: + """ + Compute the DMD loss (eq 7 in https://arxiv.org/abs/2311.18828). + Input: + - image_or_video: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images. + - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings). + - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings). + - gradient_mask: a boolean tensor with the same shape as image_or_video indicating which pixels to compute loss . + Output: + - dmd_loss: a scalar tensor representing the DMD loss. + - dmd_log_dict: a dictionary containing the intermediate tensors for logging. + """ + context_frames = _get_i2v_context_frames(image_or_video, initial_latent) + original_latent = _overwrite_i2v_context( + image_or_video, initial_latent, context_frames + ) + if clean_x is not None: + clean_x = _overwrite_i2v_context(clean_x, initial_latent, context_frames) + + batch_size, num_frame = image_or_video.shape[:2] + + with torch.no_grad(): + # Step 1: Randomly sample timestep based on the given schedule and corresponding noise + min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep + max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep + timestep = self._get_timestep( + min_timestep, + max_timestep, + batch_size, + num_frame, + self.num_frame_per_block, + # t2v keeps the original NVFP4 behavior (one shared timestep across + # all frames); i2v uses per-block timesteps. + uniform_timestep=not getattr(self.args, "i2v", False) + ) + + # TODO:should we change it to `timestep = self.scheduler.timesteps[timestep]`? + if self.timestep_shift > 1: + timestep = self.timestep_shift * \ + (timestep / 1000) / \ + (1 + (self.timestep_shift - 1) * (timestep / 1000)) * 1000 + timestep = timestep.clamp(self.min_step, self.max_step) + timestep = _zero_i2v_context_timestep(timestep, context_frames) + + noise = torch.randn_like(image_or_video) + if context_frames > 0: + noise[:, :context_frames] = 0 + noisy_latent = self.scheduler.add_noise( + original_latent.flatten(0, 1), + noise.flatten(0, 1), + timestep.flatten(0, 1) + ).detach().unflatten(0, (batch_size, num_frame)) + noisy_latent = _overwrite_i2v_context( + noisy_latent, initial_latent, context_frames + ) + + # Step 2: Compute the KL grad + grad, dmd_log_dict = self._compute_kl_grad( + noisy_image_or_video=noisy_latent, + estimated_clean_image_or_video=original_latent, + timestep=timestep, + conditional_dict=conditional_dict, + unconditional_dict=unconditional_dict, + clean_x=clean_x + ) + + context_mask = _i2v_loss_mask_like(original_latent, context_frames) + if context_mask is not None: + gradient_mask = context_mask if gradient_mask is None else gradient_mask & context_mask + + if gradient_mask is not None: + dmd_loss = 0.5 * F.mse_loss(original_latent.double( + )[gradient_mask], (original_latent.double() - grad.double()).detach()[gradient_mask], reduction="mean") + else: + dmd_loss = 0.5 * F.mse_loss(original_latent.double( + ), (original_latent.double() - grad.double()).detach(), reduction="mean") + return dmd_loss, dmd_log_dict + + def generator_loss( + self, + image_or_video_shape, + conditional_dict: dict, + unconditional_dict: dict, + clean_latent: torch.Tensor, + initial_latent: torch.Tensor = None + ) -> Tuple[torch.Tensor, dict]: + """ + Generate image/videos from noise and compute the DMD loss. + The noisy input to the generator is backward simulated. + This removes the need of any datasets during distillation. + See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details. + Input: + - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W]. + - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings). + - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings). + - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used. + Output: + - loss: a scalar tensor representing the generator loss. + - generator_log_dict: a dictionary containing the intermediate tensors for logging. + """ + # Step 1: Unroll generator to obtain fake videos + slice_last_frames = getattr(self.args, "slice_last_frames", 21) + _t_gen_start = time.time() + num_gen_frames = image_or_video_shape[1] + sampled_noise = torch.randn( + [image_or_video_shape[0], num_gen_frames, *image_or_video_shape[2:]], + device=self.device, dtype=self.dtype) + pred_image, gradient_mask, denoised_timestep_from, denoised_timestep_to = self._run_generator( + image_or_video_shape=image_or_video_shape, + conditional_dict=conditional_dict, + initial_latent=initial_latent, + slice_last_frames=slice_last_frames, + noise=sampled_noise, + clean_latent=clean_latent + ) + gen_time = time.time() - _t_gen_start + # Step 2: Compute the DMD loss + _t_loss_start = time.time() + if getattr(self.args, "teacher_forcing", False): + if getattr(self.args, "backward_simulation", True): + score_clean_x = pred_image.detach() + else: + score_clean_x = clean_latent + else: + score_clean_x = None + _bs = pred_image.shape[0] + _new_segs = pred_image.shape[1] // self.num_frame_per_block + if not getattr(self.args, "generator_is_causal", True): + _new_segs = 1 + conditional_dict = self._slice_block_cond_dict(conditional_dict, _bs, _new_segs) + unconditional_dict = self._slice_block_cond_dict(unconditional_dict, _bs, _new_segs) + dmd_loss, dmd_log_dict = self.compute_distribution_matching_loss( + image_or_video=pred_image, + conditional_dict=conditional_dict, + unconditional_dict=unconditional_dict, + gradient_mask=gradient_mask, + denoised_timestep_from=denoised_timestep_from, + denoised_timestep_to=denoised_timestep_to, + clean_x=score_clean_x, + initial_latent=initial_latent if pred_image.shape[1] == image_or_video_shape[1] else None, + ) + try: + loss_val = dmd_loss.item() + except Exception: + loss_val = float('nan') + loss_time = time.time() - _t_loss_start + + dmd_log_dict.update({ + "gen_time": gen_time, + "loss_time": loss_time + }) + + return dmd_loss, dmd_log_dict + + def critic_loss( + self, + image_or_video_shape, + conditional_dict: dict, + unconditional_dict: dict, + clean_latent: torch.Tensor, + initial_latent: torch.Tensor = None + ) -> Tuple[torch.Tensor, dict]: + """ + Generate image/videos from noise and train the critic with generated samples. + The noisy input to the generator is backward simulated. + This removes the need of any datasets during distillation. + See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details. + Input: + - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W]. + - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings). + - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings). + - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used. + Output: + - loss: a scalar tensor representing the generator loss. + - critic_log_dict: a dictionary containing the intermediate tensors for logging. + """ + slice_last_frames = getattr(self.args, "slice_last_frames", 21) + # Step 1: Run generator on backward simulated noisy input + _t_gen_start = time.time() + with torch.no_grad(): + num_gen_frames = image_or_video_shape[1] + sampled_noise = torch.randn( + [image_or_video_shape[0], num_gen_frames, *image_or_video_shape[2:]], + device=self.device, dtype=self.dtype) + generated_image, _, denoised_timestep_from, denoised_timestep_to = self._run_generator( + image_or_video_shape=image_or_video_shape, + conditional_dict=conditional_dict, + initial_latent=initial_latent, + slice_last_frames=slice_last_frames, + noise=sampled_noise, + clean_latent=clean_latent + ) + gen_time = time.time() - _t_gen_start + score_initial_latent = ( + initial_latent + if initial_latent is not None and generated_image.shape[1] == image_or_video_shape[1] + else None + ) + context_frames = _get_i2v_context_frames(generated_image, score_initial_latent) + batch_size, num_frame = generated_image.shape[:2] + + _new_segs = num_frame // self.num_frame_per_block + if not getattr(self.args, "generator_is_causal", True): + _new_segs = 1 + conditional_dict = self._slice_block_cond_dict(conditional_dict, batch_size, _new_segs) + + if getattr(self.args, "teacher_forcing", False): + if getattr(self.args, "backward_simulation", True): + score_clean_x = generated_image + else: + score_clean_x = clean_latent + else: + score_clean_x = None + if score_clean_x is not None: + score_clean_x = _overwrite_i2v_context( + score_clean_x, score_initial_latent, context_frames + ) + _t_loss_start = time.time() + + # Step 2: Compute the fake prediction + min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep + max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep + critic_timestep = self._get_timestep( + min_timestep, + max_timestep, + batch_size, + num_frame, + self.num_frame_per_block, + # t2v keeps the original NVFP4 behavior (one shared timestep across + # all frames); i2v uses per-block timesteps. + uniform_timestep=not getattr(self.args, "i2v", False) + ) + + if self.timestep_shift > 1: + critic_timestep = self.timestep_shift * \ + (critic_timestep / 1000) / (1 + (self.timestep_shift - 1) * (critic_timestep / 1000)) * 1000 + + critic_timestep = critic_timestep.clamp(self.min_step, self.max_step) + critic_timestep = _zero_i2v_context_timestep(critic_timestep, context_frames) + + critic_noise = torch.randn_like(generated_image) + if context_frames > 0: + critic_noise[:, :context_frames] = 0 + noisy_generated_image = self.scheduler.add_noise( + generated_image.flatten(0, 1), + critic_noise.flatten(0, 1), + critic_timestep.flatten(0, 1) + ).unflatten(0, (batch_size, num_frame)) + noisy_generated_image = _overwrite_i2v_context( + noisy_generated_image, score_initial_latent, context_frames + ) + + _, pred_fake_image = self.fake_score( + noisy_image_or_video=noisy_generated_image, + conditional_dict=conditional_dict, + timestep=critic_timestep, + clean_x=score_clean_x + ) + + # Step 3: Compute the denoising loss for the fake critic + if getattr(self.args, "denoising_loss_type", "flow") == "flow": + from utils.wan_5b_wrapper import WanDiffusionWrapper + flow_pred = WanDiffusionWrapper._convert_x0_to_flow_pred( + scheduler=self.scheduler, + x0_pred=pred_fake_image.flatten(0, 1), + xt=noisy_generated_image.flatten(0, 1), + timestep=critic_timestep.flatten(0, 1) + ) + pred_fake_noise = None + else: + flow_pred = None + pred_fake_noise = self.scheduler.convert_x0_to_noise( + x0=pred_fake_image.flatten(0, 1), + xt=noisy_generated_image.flatten(0, 1), + timestep=critic_timestep.flatten(0, 1) + ).unflatten(0, (batch_size, num_frame)) + + denoising_loss = self.denoising_loss_func( + x=generated_image.flatten(0, 1), + x_pred=pred_fake_image.flatten(0, 1), + noise=critic_noise.flatten(0, 1), + noise_pred=pred_fake_noise, + alphas_cumprod=self.scheduler.alphas_cumprod, + timestep=critic_timestep.flatten(0, 1), + gradient_mask=( + _i2v_loss_mask_like(generated_image, context_frames).flatten(0, 1) + if context_frames > 0 else None + ), + flow_pred=flow_pred, + ) + + try: + loss_val = denoising_loss.item() + except Exception: + loss_val = float('nan') + loss_time = time.time() - _t_loss_start + # Step 5: Debugging Log + critic_log_dict = { + "critic_timestep": critic_timestep.detach(), + "gen_time": gen_time, + "loss_time": loss_time + } + + return denoising_loss, critic_log_dict diff --git a/pipeline/__init__.py b/pipeline/__init__.py new file mode 100644 index 0000000..5e98036 --- /dev/null +++ b/pipeline/__init__.py @@ -0,0 +1,7 @@ +from .causal_diffusion_inference import CausalDiffusionInferencePipeline +from .self_forcing_training import SelfForcingTrainingPipeline + +__all__ = [ + "CausalDiffusionInferencePipeline", + "SelfForcingTrainingPipeline", +] diff --git a/pipeline/causal_diffusion_inference.py b/pipeline/causal_diffusion_inference.py new file mode 100644 index 0000000..f50eef6 --- /dev/null +++ b/pipeline/causal_diffusion_inference.py @@ -0,0 +1,1045 @@ +# Adopted from https://github.com/guandeh17/Self-Forcing +# SPDX-License-Identifier: Apache-2.0 + +from tqdm import tqdm +from typing import List, Optional +import os +import statistics +import threading +import torch +import math + +_LLV2_TIME = os.environ.get("LLV2_TIME") == "1" +_LLV2_DUMP_LATENT_DIR = os.environ.get("LLV2_DUMP_LATENT_DIR", "").strip() +# LLV2_PROFILE format: ":::", e.g. "0:20:2:2" +_LLV2_PROFILE_SPEC = os.environ.get("LLV2_PROFILE", "").strip() +_LLV2_PROFILE_OUTPUT_DIR = os.environ.get("LLV2_PROFILE_OUTPUT_DIR", "").strip() +_LLV2_PROFILE_CALL_COUNTER = 0 +from wan_5b.utils.fm_solvers import FlowDPMSolverMultistepScheduler, get_sampling_sigmas, retrieve_timesteps +from wan_5b.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler +from utils.wan_5b_wrapper import WanDiffusionWrapper, WanTextEncoder, build_vae_5b +from utils.dataset import DEFAULT_SCENE_CUT_PREFIX +from utils.config import section_get, wan_default_config +from utils.i2v_conditioning import ( + _overwrite_i2v_context, + _zero_i2v_context_timestep, +) + + +class CausalDiffusionInferencePipeline(torch.nn.Module): + def __init__( + self, + args, + device, + generator=None, + text_encoder=None, + vae=None + ): + super().__init__() + # Step 1: Initialize all models + model_name = getattr(args.model_kwargs, "model_name", "Wan2.2-TI2V-5B") + if "5B" not in model_name: + raise ValueError(f"Only Wan2.2-TI2V-5B is supported in this release, got {model_name}") + self.generator = WanDiffusionWrapper( + **getattr(args, "model_kwargs", {}), is_causal=True) if generator is None else generator + self.text_encoder = WanTextEncoder() if text_encoder is None else text_encoder + self.vae = build_vae_5b(args) if vae is None else vae + + # iter-33: optionally compile the VAE decoder (cuda:2). The Python + # `for step in range(total_steps)` wrapper in cached_decode keeps the + # mutating feat_cache in eager; per-step self.decoder() call hits the + # compiled version. Opt-in via LLV2_COMPILE_VAE=1; default off. + if os.environ.get("LLV2_COMPILE_VAE", "0") == "1": + try: + inner = getattr(self.vae, "model", None) + if inner is not None and hasattr(inner, "decoder"): + inner.decoder = torch.compile( + inner.decoder, + backend="inductor", + mode="max-autotune-no-cudagraphs", + fullgraph=False, + dynamic=False, + ) + print("[torch.compile] VAE.decoder wrapped") + except Exception as exc: + print(f"[torch.compile][warn] VAE compile setup failed: {exc}") + + # Step 2: Initialize scheduler + self.num_train_timesteps = getattr(args, "num_train_timestep", 1000) + self.sampling_steps = section_get(args, "inference", "sampling_steps", 50) + self.sample_solver = 'unipc' + self.shift = getattr(args, "timestep_shift", + getattr(args.model_kwargs, "timestep_shift", 5.0)) + + + self.frame_seq_length = math.prod(args.image_or_video_shape[-2:]) // 4 + self.model_name = model_name + self.num_transformer_blocks = wan_default_config[self.model_name]["num_transformer_blocks"] + + self.kv_cache_pos = None + self.kv_cache_neg = None + self.crossattn_cache_pos = None + self.crossattn_cache_neg = None + self.args = args + self.num_frame_per_block = getattr(args, "num_frame_per_block", 1) + self.quantize_kv = getattr(args, "kv_quant", False) + self.kv_quant_scale_rule = getattr(args, "kv_quant_scale_rule", "mse") + self.kv_quant_backend = getattr(args, "kv_quant_backend", "cuda") + self.independent_first_frame = section_get(args, "inference", "independent_first_frame", False) + self.local_attn_size = section_get( + args, "inference", "local_attn_size", -1, aliases=("inference_local_attn_size",) + ) + if self.local_attn_size == -1: + self.local_attn_size = getattr(args, "model_kwargs", {}).get("local_attn_size", -1) + self.sink_size = section_get( + args, "inference", "sink_size", None, aliases=("inference_sink_size",) + ) + if self.sink_size is None: + _model_sink = getattr(args, "model_kwargs", {}).get("sink_size", None) + if _model_sink is not None: + self.sink_size = _model_sink + if self.sink_size is None: + self.sink_size = 0 + self.scene_cut_prefix = section_get(args, "inference", "scene_cut_prefix", DEFAULT_SCENE_CUT_PREFIX) + self.multi_shot_sink = section_get(args, "inference", "multi_shot_sink", False) + self.shot_clean_recache = section_get(args, "inference", "shot_clean_recache", False) + self.global_sink_size = self.sink_size if self.multi_shot_sink else 0 + self.multi_shot_rope_offset = section_get( + args, + "inference", + "multi_shot_rope_offset", + 0.0, + ) + self.guidance_scale = section_get(args, "inference", "guidance_scale", getattr(args, "guidance_scale", 1.0)) + self.negative_prompt = section_get(args, "inference", "negative_prompt", getattr(args, "negative_prompt", "")) + self.streaming_vae = section_get(args, "inference", "streaming_vae", getattr(args, "streaming_vae", False)) + self.async_vae = section_get(args, "inference", "async_vae", getattr(args, "async_vae", False)) + vae_device = section_get(args, "inference", "vae_device", getattr(args, "vae_device", None)) + self.vae_device = torch.device(vae_device) if vae_device else None + + if self.quantize_kv: + from utils.quant import LongLiveQuantizationConfig + + self.kv_quant_config = LongLiveQuantizationConfig( + scale_rule=self.kv_quant_scale_rule, + backend=self.kv_quant_backend, + type="kv", + ) + else: + self.kv_quant_config = None + self._dit_model.kv_quant_config = self.kv_quant_config + + if self.streaming_vae and self.vae_device is not None: + vae_mode = "streaming-pipeline" + elif self.streaming_vae and self.async_vae: + vae_mode = "streaming-async" + elif self.streaming_vae: + vae_mode = "streaming" + else: + vae_mode = "batch" + print( + f"KV inference with {self.num_frame_per_block} frames per block " + f"(kv_quant={self.quantize_kv}, vae_decode={vae_mode})" + ) + + if self.num_frame_per_block > 1: + self.generator.model.num_frame_per_block = self.num_frame_per_block + self.inference_t_scale = getattr(args, "inference_t_scale", None) + self.use_relative_rope = getattr(args, "use_relative_rope", False) + self._rope_method_override = getattr(args, "rope_method", None) + self._original_seq_len_override = getattr(args, "original_seq_len", None) + + @property + def _dit_model(self): + """Return the underlying CausalWanModel, unwrapping PeftModel if present. + + After LoRA wrapping, ``self.generator.model`` is a PeftModel whose + structure is PeftModel -> LoraModel (.base_model) -> CausalWanModel + (.model). Direct attribute writes on PeftModel do NOT propagate to + CausalWanModel, so any runtime overrides (t_scale, rope_method, …) + must target the unwrapped model returned by this property. + """ + model = self.generator.model + if hasattr(model, 'base_model') and hasattr(model.base_model, 'model'): + return model.base_model.model + return model + + def inference( + self, + noise: torch.Tensor, + text_prompts: List[str], + initial_latent: Optional[torch.Tensor] = None, + return_latents: bool = False, + start_frame_index: Optional[int] = 0 + ) -> torch.Tensor: + """ + Perform inference on the given noise and text prompts. + Inputs: + noise (torch.Tensor): The input noise tensor of shape + (batch_size, num_output_frames, num_channels, height, width). + text_prompts (List[str]): The list of text prompts. + initial_latent (torch.Tensor): The initial latent tensor of shape + (batch_size, num_input_frames, num_channels, height, width). + If num_input_frames is 1, perform image to video. + If num_input_frames is greater than 1, perform video extension. + return_latents (bool): Whether to return the latents. + start_frame_index (int): In long video generation, where does the current window start? + Outputs: + video (torch.Tensor): The generated video tensor of shape + (batch_size, num_frames, num_channels, height, width). It is normalized to be in the range [0, 1]. + """ + batch_size, num_frames, num_channels, height, width = noise.shape + num_input_frames = initial_latent.shape[1] if initial_latent is not None else 0 + clamp_i2v_first_chunk = self.independent_first_frame and initial_latent is not None + if clamp_i2v_first_chunk and num_input_frames != 1: + raise ValueError( + f"i2v first-chunk clamp expects one conditioning latent frame, got {num_input_frames}." + ) + + if not self.independent_first_frame or clamp_i2v_first_chunk: + # If the first frame is independent and the first frame is provided, then the number of frames in the + # noise should still be a multiple of num_frame_per_block + assert num_frames % self.num_frame_per_block == 0 + num_blocks = num_frames // self.num_frame_per_block + elif self.independent_first_frame and initial_latent is None: + # Using a [1, 4, 4, 4, 4, 4] model to generate a video without image conditioning + assert (num_frames - 1) % self.num_frame_per_block == 0 + num_blocks = (num_frames - 1) // self.num_frame_per_block + num_output_frames = ( + num_frames if clamp_i2v_first_chunk else num_frames + num_input_frames + ) + conditional_dict = self.text_encoder( + text_prompts=text_prompts[0] + ) + conditional_dict_list = [ + {"prompt_embeds": conditional_dict["prompt_embeds"][i:i+1]} + for i in range(conditional_dict["prompt_embeds"].shape[0]) + ] + use_cfg = self.guidance_scale != 1.0 + if use_cfg: + unconditional_dict = self.text_encoder( + text_prompts=[self.negative_prompt] * batch_size + ) + else: + unconditional_dict = None + + output = torch.zeros( + [batch_size, num_output_frames, num_channels, height, width], + device=noise.device, + dtype=noise.dtype + ) + + # Step 1: Initialize KV cache to all zeros + if self.kv_cache_pos is None: + self._initialize_kv_cache( + batch_size=batch_size, + dtype=noise.dtype, + device=noise.device + ) + self._initialize_crossattn_cache( + batch_size=batch_size, + dtype=noise.dtype, + device=noise.device + ) + else: + # reset cross attn cache + for block_index in range(self.num_transformer_blocks): + self.crossattn_cache_pos[block_index]["is_init"] = False + if use_cfg: + self.crossattn_cache_neg[block_index]["is_init"] = False + # reset kv cache + for block_index in range(len(self.kv_cache_pos)): + self.kv_cache_pos[block_index]["global_end_index"] = torch.tensor( + [0], dtype=torch.long, device=noise.device) + self.kv_cache_pos[block_index]["local_end_index"] = torch.tensor( + [0], dtype=torch.long, device=noise.device) + self.kv_cache_pos[block_index]["pinned_start"].fill_(-1) + self.kv_cache_pos[block_index]["pinned_len"].zero_() + if use_cfg: + self.kv_cache_neg[block_index]["global_end_index"] = torch.tensor( + [0], dtype=torch.long, device=noise.device) + self.kv_cache_neg[block_index]["local_end_index"] = torch.tensor( + [0], dtype=torch.long, device=noise.device) + self.kv_cache_neg[block_index]["pinned_start"].fill_(-1) + self.kv_cache_neg[block_index]["pinned_len"].zero_() + + # Step 2: Cache context feature + current_start_frame = start_frame_index + cache_start_frame = 0 + + # Save model state before overriding for inference. + # Use _dit_model to reach the real CausalWanModel (PeftModel wrapping + # intercepts attribute writes, so self.generator.model.xxx would land + # on the wrapper instead of the model that reads them in forward()). + dit = self._dit_model + prev_local_attn_size = dit.local_attn_size + prev_t_scale = getattr(dit, 't_scale', 1.0) + prev_rope_method = getattr(dit, 'rope_method', 'linear') + prev_original_seq_len = getattr(dit, 'original_seq_len', None) + prev_use_relative_rope = getattr(dit, 'use_relative_rope', False) + prev_rope_temporal_offset = getattr(dit, 'rope_temporal_offset', 0.0) + prev_max_attention_sizes = {} + prev_sink_sizes = {} + prev_global_sink_sizes = {} + for name, module in self.generator.model.named_modules(): + if hasattr(module, 'max_attention_size'): + prev_max_attention_sizes[name] = module.max_attention_size + if hasattr(module, 'sink_size'): + prev_sink_sizes[name] = module.sink_size + if hasattr(module, 'global_sink_size'): + prev_global_sink_sizes[name] = module.global_sink_size + + dit.local_attn_size = self.local_attn_size + print(f"[inference] local_attn_size set on model: {dit.local_attn_size}") + self._set_all_modules_max_attention_size(self.local_attn_size) + + if self.sink_size is not None: + self._set_all_modules_sink_size(self.sink_size) + print(f"[inference] sink_size set to: {self.sink_size}" + f"{', multi_shot_sink enabled (pinned position)' if self.multi_shot_sink else ''}" + f"{', shot_clean_recache enabled' if self.shot_clean_recache else ''}") + + # Propagate the internally derived global sink length. + self._set_all_modules_global_sink_size(self.global_sink_size) + if self.global_sink_size and self.global_sink_size > 0: + print(f"[inference] auto_global_sink_size set to: {self.global_sink_size} " + f"(first {self.global_sink_size} frames permanently anchored)") + + if self.inference_t_scale is not None: + dit.t_scale = self.inference_t_scale + print(f"[inference] t_scale overridden to: {dit.t_scale}") + + if self._rope_method_override is not None: + dit.rope_method = self._rope_method_override + if self._original_seq_len_override is not None: + dit.original_seq_len = self._original_seq_len_override + print(f"[inference] rope_method={dit.rope_method}, " + f"original_seq_len={dit.original_seq_len}") + + dit.use_relative_rope = self.use_relative_rope + if self.use_relative_rope: + print(f"[inference] use_relative_rope enabled") + + dit.rope_temporal_offset = 0.0 + if self.multi_shot_rope_offset != 0.0: + print(f"[inference] multi_shot_rope_offset={self.multi_shot_rope_offset} " + f"(multi-shot RoPE offset enabled)") + + try: + raw_prompts = text_prompts[0] if isinstance(text_prompts[0], (list, tuple)) else text_prompts + return self._inference_inner( + noise=noise, batch_size=batch_size, num_frames=num_frames, + num_channels=num_channels, height=height, width=width, + num_blocks=num_blocks, num_input_frames=num_input_frames, + num_output_frames=num_output_frames, output=output, + conditional_dict=conditional_dict, + conditional_dict_list=conditional_dict_list, + unconditional_dict=unconditional_dict, + use_cfg=use_cfg, initial_latent=initial_latent, + clamp_i2v_first_chunk=clamp_i2v_first_chunk, + return_latents=return_latents, + current_start_frame=current_start_frame, + cache_start_frame=cache_start_frame, + raw_prompts=raw_prompts, + ) + finally: + dit.local_attn_size = prev_local_attn_size + dit.t_scale = prev_t_scale + dit.rope_method = prev_rope_method + dit.original_seq_len = prev_original_seq_len + dit.use_relative_rope = prev_use_relative_rope + dit.rope_temporal_offset = prev_rope_temporal_offset + for name, module in self.generator.model.named_modules(): + if name in prev_max_attention_sizes: + try: + module.max_attention_size = prev_max_attention_sizes[name] + except Exception: + pass + if name in prev_sink_sizes: + try: + module.sink_size = prev_sink_sizes[name] + except Exception: + pass + if name in prev_global_sink_sizes: + try: + module.global_sink_size = prev_global_sink_sizes[name] + except Exception: + pass + + def _inference_inner( + self, noise, batch_size, num_frames, num_channels, height, width, + num_blocks, num_input_frames, num_output_frames, output, + conditional_dict, conditional_dict_list, unconditional_dict, + use_cfg, initial_latent, clamp_i2v_first_chunk, return_latents, + current_start_frame, cache_start_frame, + raw_prompts=None, + ): + + if initial_latent is not None and not clamp_i2v_first_chunk: + timestep = torch.ones([batch_size, 1], device=noise.device, dtype=torch.int64) * 0 + if self.independent_first_frame: + # Assume num_input_frames is 1 + self.num_frame_per_block * num_input_blocks + assert (num_input_frames - 1) % self.num_frame_per_block == 0 + num_input_blocks = (num_input_frames - 1) // self.num_frame_per_block + output[:, :1] = initial_latent[:, :1] + self.generator( + noisy_image_or_video=initial_latent[:, :1], + conditional_dict=conditional_dict, + timestep=timestep * 0, + kv_cache=self.kv_cache_pos, + crossattn_cache=self.crossattn_cache_pos, + current_start=current_start_frame * self.frame_seq_length, + cache_start=cache_start_frame * self.frame_seq_length + ) + if use_cfg: + self.generator( + noisy_image_or_video=initial_latent[:, :1], + conditional_dict=unconditional_dict, + timestep=timestep * 0, + kv_cache=self.kv_cache_neg, + crossattn_cache=self.crossattn_cache_neg, + current_start=current_start_frame * self.frame_seq_length, + cache_start=cache_start_frame * self.frame_seq_length + ) + current_start_frame += 1 + cache_start_frame += 1 + else: + # Assume num_input_frames is self.num_frame_per_block * num_input_blocks + assert num_input_frames % self.num_frame_per_block == 0 + num_input_blocks = num_input_frames // self.num_frame_per_block + + for block_index in range(num_input_blocks): + current_ref_latents = \ + initial_latent[:, cache_start_frame:cache_start_frame + self.num_frame_per_block] + output[:, cache_start_frame:cache_start_frame + self.num_frame_per_block] = current_ref_latents + self.generator( + noisy_image_or_video=current_ref_latents, + conditional_dict=conditional_dict, + timestep=timestep * 0, + kv_cache=self.kv_cache_pos, + crossattn_cache=self.crossattn_cache_pos, + current_start=current_start_frame * self.frame_seq_length, + cache_start=cache_start_frame * self.frame_seq_length + ) + if use_cfg: + self.generator( + noisy_image_or_video=current_ref_latents, + conditional_dict=unconditional_dict, + timestep=timestep * 0, + kv_cache=self.kv_cache_neg, + crossattn_cache=self.crossattn_cache_neg, + current_start=current_start_frame * self.frame_seq_length, + cache_start=cache_start_frame * self.frame_seq_length + ) + current_start_frame += self.num_frame_per_block + cache_start_frame += self.num_frame_per_block + + # Step 3: Temporal denoising loop + all_num_frames = [self.num_frame_per_block] * num_blocks + if self.independent_first_frame and initial_latent is None: + all_num_frames = [1] + all_num_frames + + # Multi-shot RoPE offset: track current shot index for phase offset. + current_shot_index = 0 + phi = self.multi_shot_rope_offset + self._dit_model.rope_temporal_offset = 0.0 + streaming_decode = self.streaming_vae and not return_latents + pipeline_vae = streaming_decode and self.vae_device is not None + async_vae = streaming_decode and self.async_vae and not pipeline_vae + if streaming_decode: + vae_dev = self.vae_device if pipeline_vae else noise.device + vae_scale = [ + self.vae.mean.to(device=vae_dev, dtype=noise.dtype), + 1.0 / self.vae.std.to(device=vae_dev, dtype=noise.dtype), + ] + self.vae.model.clear_cache() + video_chunks = [] + if async_vae: + vae_stream = torch.cuda.Stream(device=noise.device) + prev_vae_done = None + if pipeline_vae: + vae_thread_error = [] + vae_thread_chunks = [] + vae_work_queue = [] + vae_queue_lock = threading.Lock() + vae_work_ready = threading.Event() + vae_all_done = threading.Event() + + def _vae_thread_fn(): + try: + while True: + vae_work_ready.wait() + vae_work_ready.clear() + while True: + with vae_queue_lock: + if not vae_work_queue: + break + item = vae_work_queue.pop(0) + if item is None: + vae_all_done.set() + return + decoded = self.vae.model.cached_decode( + item, + vae_scale, + ).float().clamp_(-1, 1) + # Pinned-memory DtoH: pageable copy hits ~0.2 GB/s + # (1.5s per 313MB chunk → ~80s/prompt at end); pinned + # path runs at PCIe limit (~25 GB/s = ~12ms / chunk). + pinned = torch.empty( + decoded.shape, dtype=decoded.dtype, + device="cpu", pin_memory=True, + ) + pinned.copy_(decoded, non_blocking=True) + torch.cuda.synchronize(decoded.device) + vae_thread_chunks.append(pinned) + except Exception as exc: + vae_thread_error.append(exc) + vae_all_done.set() + + vae_bg_thread = threading.Thread(target=_vae_thread_fn, daemon=True) + vae_bg_thread.start() + + _block_events = [] if _LLV2_TIME else None + global _LLV2_PROFILE_CALL_COUNTER + _call_idx = _LLV2_PROFILE_CALL_COUNTER + _LLV2_PROFILE_CALL_COUNTER += 1 + _prof = None + _prof_trace_path = None + if _LLV2_PROFILE_SPEC and _LLV2_PROFILE_OUTPUT_DIR: + _parts = _LLV2_PROFILE_SPEC.split(":") + _target_call = int(_parts[0]) if len(_parts) > 0 else 0 + _wait_n = int(_parts[1]) if len(_parts) > 1 else 20 + _warmup_n = int(_parts[2]) if len(_parts) > 2 else 2 + _active_n = int(_parts[3]) if len(_parts) > 3 else 2 + if _call_idx == _target_call: + from torch.profiler import profile as _tp_profile, ProfilerActivity, schedule + os.makedirs(_LLV2_PROFILE_OUTPUT_DIR, exist_ok=True) + _prof_trace_path = os.path.join( + _LLV2_PROFILE_OUTPUT_DIR, f"trace_call{_call_idx}.json" + ) + _prof = _tp_profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + schedule=schedule( + wait=_wait_n, warmup=_warmup_n, active=_active_n, repeat=1 + ), + record_shapes=False, + with_stack=False, + ) + _prof.start() + print( + f"[LLV2_PROFILE] enabled for call={_call_idx} " + f"wait={_wait_n} warmup={_warmup_n} active={_active_n} " + f"-> {_prof_trace_path}", + flush=True, + ) + for chunk_index, current_num_frames in enumerate(all_num_frames): + if _LLV2_TIME: + _ev_s = torch.cuda.Event(enable_timing=True) + _ev_e = torch.cuda.Event(enable_timing=True) + _ev_s.record() + conditional_dict = conditional_dict_list[chunk_index] + # Reset the cross-attention cache when each chunk uses a different + # prompt; otherwise the model reuses the previous chunk's k/v and + # ignores the current conditional_dict. + for block_index in range(self.num_transformer_blocks): + self.crossattn_cache_pos[block_index]["is_init"] = False + self.crossattn_cache_neg[block_index]["is_init"] = False + + # Update RoPE phase offset on shot boundaries. + is_shot_boundary = self._is_shot_boundary(raw_prompts, chunk_index) + if is_shot_boundary and phi != 0.0: + current_shot_index += 1 + self._dit_model.rope_temporal_offset = current_shot_index * phi + print(f"[inference] multi-shot RoPE: shot_index={current_shot_index}, " + f"temporal_offset={self._dit_model.rope_temporal_offset:.4f}") + + first_i2v_block = clamp_i2v_first_chunk and chunk_index == 0 + noise_start_frame = ( + cache_start_frame + if clamp_i2v_first_chunk + else cache_start_frame - num_input_frames + ) + noisy_input = noise[ + :, + noise_start_frame:noise_start_frame + current_num_frames, + ] + latents = noisy_input + + # Step 3.1: Spatial denoising loop + sample_scheduler = self._initialize_sample_scheduler(noise) + for _, t in enumerate(tqdm(sample_scheduler.timesteps)): + timestep = t * torch.ones( + [batch_size, current_num_frames], device=noise.device, dtype=torch.float32 + ) + if first_i2v_block: + latents = _overwrite_i2v_context( + latents, initial_latent, num_input_frames + ) + timestep = _zero_i2v_context_timestep( + timestep, num_input_frames + ) + latent_model_input = latents + + flow_pred_cond, _ = self.generator( + noisy_image_or_video=latent_model_input, + conditional_dict=conditional_dict, + timestep=timestep, + kv_cache=self.kv_cache_pos, + crossattn_cache=self.crossattn_cache_pos, + current_start=current_start_frame * self.frame_seq_length, + cache_start=cache_start_frame * self.frame_seq_length + ) + if use_cfg: + flow_pred_uncond, _ = self.generator( + noisy_image_or_video=latent_model_input, + conditional_dict=unconditional_dict, + timestep=timestep, + kv_cache=self.kv_cache_neg, + crossattn_cache=self.crossattn_cache_neg, + current_start=current_start_frame * self.frame_seq_length, + cache_start=cache_start_frame * self.frame_seq_length + ) + flow_pred = flow_pred_uncond + self.guidance_scale * ( + flow_pred_cond - flow_pred_uncond) + else: + flow_pred = flow_pred_cond + + temp_x0 = sample_scheduler.step( + flow_pred, + t, + latents, + return_dict=False)[0] + latents = temp_x0 + if first_i2v_block: + latents = _overwrite_i2v_context( + latents, initial_latent, num_input_frames + ) + # iter-34: removed per-step debug print of kv_cache scalar + # tensors (was forcing GPU→CPU sync every sampling step × + # 4 steps × 48 chunks = 192 stalls per prompt). Re-enable + # behind LLV2_DEBUG_KV=1 if needed for debugging. + if os.environ.get("LLV2_DEBUG_KV", "0") == "1": + print(f"kv_cache['local_end_index']: {self.kv_cache_pos[0]['local_end_index']}") + print(f"kv_cache['global_end_index']: {self.kv_cache_pos[0]['global_end_index']}") + + # Step 3.2: record the model's output + if first_i2v_block: + latents = _overwrite_i2v_context( + latents, initial_latent, num_input_frames + ) + output[:, cache_start_frame:cache_start_frame + current_num_frames] = latents + + # Step 3.3: rerun with timestep zero to update KV cache using clean context + is_scene_cut = self._is_scene_cut(raw_prompts, chunk_index) + + if is_scene_cut and self.shot_clean_recache: + print(f"[inference] Scene cut at chunk {chunk_index}, zeroing KV before recache") + current_start_tokens = current_start_frame * self.frame_seq_length + self._zero_kv_data(self.kv_cache_pos, current_start_tokens) + if use_cfg: + self._zero_kv_data(self.kv_cache_neg, current_start_tokens) + + self.generator( + noisy_image_or_video=latents, + conditional_dict=conditional_dict, + timestep=timestep * 0, + kv_cache=self.kv_cache_pos, + crossattn_cache=self.crossattn_cache_pos, + current_start=current_start_frame * self.frame_seq_length, + cache_start=cache_start_frame * self.frame_seq_length + ) + if use_cfg: + self.generator( + noisy_image_or_video=latents, + conditional_dict=unconditional_dict, + timestep=timestep * 0, + kv_cache=self.kv_cache_neg, + crossattn_cache=self.crossattn_cache_neg, + current_start=current_start_frame * self.frame_seq_length, + cache_start=cache_start_frame * self.frame_seq_length + ) + + # Step 3.3b: pin the current chunk for multi-shot sink on scene cut. + if is_scene_cut: + print(f"[inference] Scene cut at chunk {chunk_index}, pinning chunk as shot-sink") + self._pin_current_chunk(self.kv_cache_pos, current_num_frames) + if use_cfg: + self._pin_current_chunk(self.kv_cache_neg, current_num_frames) + + if streaming_decode: + if async_vae: + diffusion_done = torch.cuda.Event() + diffusion_done.record() + if prev_vae_done is not None: + prev_vae_done.synchronize() + with torch.cuda.stream(vae_stream): + vae_stream.wait_event(diffusion_done) + chunk_bcthw = latents.permute(0, 2, 1, 3, 4).contiguous() + decoded_chunk = self.vae.model.cached_decode( + chunk_bcthw, + vae_scale, + ).float().clamp_(-1, 1) + video_chunks.append(decoded_chunk) + prev_vae_done = torch.cuda.Event() + prev_vae_done.record(vae_stream) + elif pipeline_vae: + latent_on_vae = latents.permute(0, 2, 1, 3, 4).contiguous().to(vae_dev) + with vae_queue_lock: + vae_work_queue.append(latent_on_vae) + vae_work_ready.set() + else: + chunk_bcthw = latents.permute(0, 2, 1, 3, 4).contiguous() + decoded_chunk = self.vae.model.cached_decode( + chunk_bcthw, + vae_scale, + ).float().clamp_(-1, 1) + video_chunks.append(decoded_chunk.cpu()) + del decoded_chunk, chunk_bcthw + torch.cuda.empty_cache() + + # Step 3.4: update the start and end frame indices + current_start_frame += current_num_frames + cache_start_frame += current_num_frames + + if _LLV2_TIME: + _ev_e.record() + _block_events.append((_ev_s, _ev_e)) + if _prof is not None: + _prof.step() + + if _prof is not None: + _prof.stop() + _prof.export_chrome_trace(_prof_trace_path) + print(f"[LLV2_PROFILE] saved trace -> {_prof_trace_path}", flush=True) + + if _LLV2_TIME and _block_events: + torch.cuda.synchronize() + _times = [_s.elapsed_time(_e) for _s, _e in _block_events] + _sorted = sorted(_times) + _n = len(_sorted) + _p10 = _sorted[max(0, _n // 10 - 1)] + _p50 = statistics.median(_sorted) + _p90 = _sorted[max(0, _n - _n // 10 - 1)] + _mean = sum(_times) / _n + _total = sum(_times) + print( + f"[LLV2_TIME] blocks={_n} mean_ms={_mean:.2f} p10_ms={_p10:.2f} " + f"median_ms={_p50:.2f} p90_ms={_p90:.2f} total_ms={_total:.2f}", + flush=True, + ) + + if _LLV2_DUMP_LATENT_DIR: + os.makedirs(_LLV2_DUMP_LATENT_DIR, exist_ok=True) + _existing = sum( + 1 for _f in os.listdir(_LLV2_DUMP_LATENT_DIR) + if _f.startswith("latent_") and _f.endswith(".pt") + ) + _path = os.path.join(_LLV2_DUMP_LATENT_DIR, f"latent_{_existing:04d}.pt") + torch.save(output.detach().cpu(), _path) + print(f"[LLV2_DUMP] saved latent {tuple(output.shape)} -> {_path}", flush=True) + + # Step 4: Decode the output + + + if return_latents: + return output + elif streaming_decode: + if async_vae: + vae_stream.synchronize() + elif pipeline_vae: + with vae_queue_lock: + vae_work_queue.append(None) + vae_work_ready.set() + vae_all_done.wait() + vae_bg_thread.join() + if vae_thread_error: + raise RuntimeError( + f"[pipeline_vae] VAE decode failed: {vae_thread_error[0]}" + ) from vae_thread_error[0] + video_chunks = vae_thread_chunks + video_bcthw = torch.cat(video_chunks, dim=2) + video = video_bcthw.permute(0, 2, 1, 3, 4) + video = (video * 0.5 + 0.5).clamp(0, 1) + self.vae.model.clear_cache() + return video + else: + video = self.vae.decode_to_pixel(output) + video = (video * 0.5 + 0.5).clamp(0, 1) + return video + + def _initialize_kv_cache(self, batch_size, dtype, device): + """ + Initialize a Per-GPU KV cache for the Wan model. + """ + kv_cache_pos = [] + kv_cache_neg = [] + num_heads = wan_default_config[self.model_name]["num_heads"] + head_dim = wan_default_config[self.model_name]["head_dim"] + if self.local_attn_size != -1: + # Use the local attention size to compute the KV cache size + kv_cache_size = self.local_attn_size * self.frame_seq_length + else: + # Use the default KV cache size + kv_cache_size = 3 * self.num_frame_per_block * self.frame_seq_length + + block_token_size = self.num_frame_per_block * self.frame_seq_length + max_blocks = kv_cache_size // block_token_size + + if self.quantize_kv: + from utils.quant import clone_quantized_tensor, quantize_to_fp4 + + print( + f"[KV Cache] Quantized (nvfp4): block_token_size={block_token_size}, " + f"max_blocks={max_blocks}, num_heads={num_heads}, layers={self.num_transformer_blocks}" + ) + zero_block = torch.zeros( + [block_token_size * num_heads, head_dim], + dtype=dtype, + device=device, + ) + zero_qt = quantize_to_fp4(zero_block, self.kv_quant_config) + + for _ in range(self.num_transformer_blocks): + if self.quantize_kv: + kv_cache_pos.append({ + "k": [clone_quantized_tensor(zero_qt) for _ in range(max_blocks)], + "v": [clone_quantized_tensor(zero_qt) for _ in range(max_blocks)], + "quantized": True, + "block_token_size": block_token_size, + "max_blocks": max_blocks, + "num_heads": num_heads, + "num_filled_blocks": 0, + "global_end_index": torch.tensor([0], dtype=torch.long, device=device), + "local_end_index": torch.tensor([0], dtype=torch.long, device=device), + "pinned_start": torch.tensor([-1], dtype=torch.long, device=device), + "pinned_len": torch.tensor([0], dtype=torch.long, device=device), + }) + kv_cache_neg.append({ + "k": [clone_quantized_tensor(zero_qt) for _ in range(max_blocks)], + "v": [clone_quantized_tensor(zero_qt) for _ in range(max_blocks)], + "quantized": True, + "block_token_size": block_token_size, + "max_blocks": max_blocks, + "num_heads": num_heads, + "num_filled_blocks": 0, + "global_end_index": torch.tensor([0], dtype=torch.long, device=device), + "local_end_index": torch.tensor([0], dtype=torch.long, device=device), + "pinned_start": torch.tensor([-1], dtype=torch.long, device=device), + "pinned_len": torch.tensor([0], dtype=torch.long, device=device), + }) + else: + kv_cache_pos.append({ + "k": torch.zeros([batch_size, kv_cache_size, num_heads, head_dim], dtype=dtype, device=device), + "v": torch.zeros([batch_size, kv_cache_size, num_heads, head_dim], dtype=dtype, device=device), + "quantized": False, + "block_token_size": block_token_size, + "max_blocks": max_blocks, + "num_heads": num_heads, + "num_filled_blocks": 0, + "global_end_index": torch.tensor([0], dtype=torch.long, device=device), + "local_end_index": torch.tensor([0], dtype=torch.long, device=device), + "pinned_start": torch.tensor([-1], dtype=torch.long, device=device), + "pinned_len": torch.tensor([0], dtype=torch.long, device=device), + }) + kv_cache_neg.append({ + "k": torch.zeros([batch_size, kv_cache_size, num_heads, head_dim], dtype=dtype, device=device), + "v": torch.zeros([batch_size, kv_cache_size, num_heads, head_dim], dtype=dtype, device=device), + "quantized": False, + "block_token_size": block_token_size, + "max_blocks": max_blocks, + "num_heads": num_heads, + "num_filled_blocks": 0, + "global_end_index": torch.tensor([0], dtype=torch.long, device=device), + "local_end_index": torch.tensor([0], dtype=torch.long, device=device), + "pinned_start": torch.tensor([-1], dtype=torch.long, device=device), + "pinned_len": torch.tensor([0], dtype=torch.long, device=device), + }) + + self.kv_cache_pos = kv_cache_pos # always store the clean cache + self.kv_cache_neg = kv_cache_neg # always store the clean cache + + def _initialize_crossattn_cache(self, batch_size, dtype, device): + """ + Initialize a Per-GPU cross-attention cache for the Wan model. + """ + crossattn_cache_pos = [] + crossattn_cache_neg = [] + num_heads = wan_default_config[self.model_name]["num_heads"] + head_dim = wan_default_config[self.model_name]["head_dim"] + for _ in range(self.num_transformer_blocks): + crossattn_cache_pos.append({ + "k": torch.zeros([batch_size, 512, num_heads, head_dim], dtype=dtype, device=device), + "v": torch.zeros([batch_size, 512, num_heads, head_dim], dtype=dtype, device=device), + "is_init": False + }) + crossattn_cache_neg.append({ + "k": torch.zeros([batch_size, 512, num_heads, head_dim], dtype=dtype, device=device), + "v": torch.zeros([batch_size, 512, num_heads, head_dim], dtype=dtype, device=device), + "is_init": False + }) + + self.crossattn_cache_pos = crossattn_cache_pos # always store the clean cache + self.crossattn_cache_neg = crossattn_cache_neg # always store the clean cache + + def clear_cache(self): + """ + Explicitly release large KV / cross-attention caches to free GPU memory. + Safe to call between independent inference calls; caches will be + re-created on demand by _initialize_kv_cache/_initialize_crossattn_cache. + """ + self.kv_cache_pos = None + self.kv_cache_neg = None + self.crossattn_cache_pos = None + self.crossattn_cache_neg = None + + def _initialize_sample_scheduler(self, noise): + if self.sample_solver == 'unipc': + sample_scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=self.num_train_timesteps, + shift=1, + use_dynamic_shifting=False) + sample_scheduler.set_timesteps( + self.sampling_steps, device=noise.device, shift=self.shift) + self.timesteps = sample_scheduler.timesteps + elif self.sample_solver == 'dpm++': + sample_scheduler = FlowDPMSolverMultistepScheduler( + num_train_timesteps=self.num_train_timesteps, + shift=1, + use_dynamic_shifting=False) + sampling_sigmas = get_sampling_sigmas(self.sampling_steps, self.shift) + self.timesteps, _ = retrieve_timesteps( + sample_scheduler, + device=noise.device, + sigmas=sampling_sigmas) + else: + raise NotImplementedError("Unsupported solver.") + return sample_scheduler + + def _set_all_modules_max_attention_size(self, local_attn_size_value: int): + """ + Set max_attention_size on all submodules that define it. + If local_attn_size_value == -1, use the model's global default (32760 for Wan, 28160 for 5B). + Otherwise, set to local_attn_size_value * frame_seq_length. + """ + if local_attn_size_value == -1: + target_size = 32760 + policy = "global" + else: + target_size = int(local_attn_size_value) * self.frame_seq_length + policy = "local" + + updated_modules = [] + # Update root model if applicable + if hasattr(self.generator.model, "max_attention_size"): + try: + prev = getattr(self.generator.model, "max_attention_size") + except Exception: + prev = None + setattr(self.generator.model, "max_attention_size", target_size) + updated_modules.append("") + + # Update all child modules + for name, module in self.generator.model.named_modules(): + if hasattr(module, "max_attention_size"): + try: + prev = getattr(module, "max_attention_size") + except Exception: + prev = None + try: + setattr(module, "max_attention_size", target_size) + updated_modules.append(name if name else module.__class__.__name__) + except Exception: + pass + + def _set_all_modules_sink_size(self, sink_size_value: int): + """ + Override sink_size on all submodules that define it. + """ + if hasattr(self.generator.model, "sink_size"): + setattr(self.generator.model, "sink_size", sink_size_value) + + for name, module in self.generator.model.named_modules(): + if hasattr(module, "sink_size"): + try: + setattr(module, "sink_size", sink_size_value) + except Exception: + pass + + def _set_all_modules_global_sink_size(self, value: int): + """Override global_sink_size on all submodules; create the attribute if missing.""" + setattr(self.generator.model, "global_sink_size", value) + for _, module in self.generator.model.named_modules(): + try: + setattr(module, "global_sink_size", value) + except Exception: + pass + + def _is_shot_boundary(self, raw_prompts, chunk_index): + """Return True when *chunk_index* starts a new shot (prompt-based detection). + + Pure prompt check — no dependency on sink config so that Narrative + RoPE and other shot-aware features can reuse it independently. + """ + if chunk_index == 0: + return False + if not isinstance(raw_prompts, (list, tuple)): + return False + if chunk_index >= len(raw_prompts): + return False + prompt = raw_prompts[chunk_index] + return isinstance(prompt, str) and prompt.startswith(self.scene_cut_prefix) + + def _is_scene_cut(self, raw_prompts, chunk_index): + """Return True when *chunk_index* is the first chunk of a new scene + AND multi-shot sink is enabled.""" + if not self.multi_shot_sink: + return False + if not self.sink_size or self.sink_size == 0: + return False + return self._is_shot_boundary(raw_prompts, chunk_index) + + def _update_sink_for_scene_cut(self, kv_cache, current_num_frames): + """Legacy copy-to-front sink relocation (used by training pipeline).""" + global_sink_tokens = self.global_sink_size * self.frame_seq_length + shot_sink_tokens = self.sink_size * self.frame_seq_length + chunk_tokens = current_num_frames * self.frame_seq_length + copy_len = min(shot_sink_tokens, chunk_tokens) + + # iter-38: local_end_index is in lockstep across all blocks (set by + # _apply_cache_updates with the same value). Read once → 60 syncs to 1. + local_end = int(kv_cache[0]["local_end_index"].item()) + chunk_start = local_end - chunk_tokens + dst_start = global_sink_tokens + src_slice = slice(chunk_start, chunk_start + copy_len) + dst_slice = slice(dst_start, dst_start + copy_len) + + for block_cache in kv_cache: + block_cache["k"][:, dst_slice] = block_cache["k"][:, src_slice].clone() + block_cache["v"][:, dst_slice] = block_cache["v"][:, src_slice].clone() + + def _pin_current_chunk(self, kv_cache, current_num_frames): + """Mark the current chunk's buffer position as pinned for multi-shot sink. + + The pinned region REPLACES the original sink on the next rolling event. + No data is copied here — relocation happens inside the attention layer + during rolling, ensuring zero duplication. + """ + chunk_tokens = current_num_frames * self.frame_seq_length + pin_len = min(self.sink_size * self.frame_seq_length, chunk_tokens) + + # iter-38: local_end_index is in lockstep across all blocks. Read once. + local_end = int(kv_cache[0]["local_end_index"].item()) + chunk_start = local_end - chunk_tokens + + for block_cache in kv_cache: + block_cache["pinned_start"].fill_(chunk_start) + block_cache["pinned_len"].fill_(pin_len) + + def _zero_kv_data(self, kv_cache, current_start_tokens): + """Reset KV cache for clean recache, preserving global sink.""" + global_sink_tokens = self.global_sink_size * self.frame_seq_length + for block_cache in kv_cache: + block_cache["local_end_index"].fill_(global_sink_tokens) + block_cache["global_end_index"].fill_(current_start_tokens) + block_cache["pinned_start"].fill_(-1) + block_cache["pinned_len"].zero_() diff --git a/pipeline/causal_diffusion_inference_sp.py b/pipeline/causal_diffusion_inference_sp.py new file mode 100644 index 0000000..3bddfdf --- /dev/null +++ b/pipeline/causal_diffusion_inference_sp.py @@ -0,0 +1,254 @@ +# Copyright 2024-2025 LongLive Authors. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Sequence-parallel causal diffusion inference pipeline for Wan2.2-TI2V-5B.""" + +from typing import List, Optional +import types + +import torch +import torch.distributed as dist + +from pipeline.causal_diffusion_inference import CausalDiffusionInferencePipeline +from utils.config import wan_default_config +from utils.scheduler import FlowMatchScheduler, SchedulerInterface + +from wan_5b.distributed.sp_ulysses_inference import ( + get_sp_world_size, + init_sequence_parallel, + is_sp_enabled, + sp_print, +) + + +def _model_kw(model_kwargs, key, default=None): + if isinstance(model_kwargs, dict): + return model_kwargs.get(key, default) + return getattr(model_kwargs, key, default) + + +class SPWanDiffusionWrapper5B(torch.nn.Module): + """Wan 5B diffusion wrapper backed by the Ulysses SP causal model.""" + + def __init__( + self, + model_name="Wan2.2-TI2V-5B", + timestep_shift=5.0, + local_attn_size=-1, + sink_size=0, + num_frame_per_block=1, + t_scale=1.0, + rope_method="linear", + original_seq_len=None, + ): + super().__init__() + if model_name != "Wan2.2-TI2V-5B": + raise ValueError(f"SP inference only supports Wan2.2-TI2V-5B, got {model_name}") + + from wan_5b.modules.causal_model_sp_ulysses import UlyssesSPCausalWanModel + + sp_print("Using Ulysses SP model for Wan2.2-TI2V-5B") + self.model = UlyssesSPCausalWanModel( + model_type="ti2v", + in_dim=48, + dim=3072, + ffn_dim=14336, + out_dim=48, + num_heads=24, + num_layers=30, + local_attn_size=local_attn_size, + sink_size=sink_size, + num_frame_per_block=num_frame_per_block, + ) + self.model.eval() + self.model.t_scale = t_scale + self.model.rope_method = rope_method + self.model.original_seq_len = original_seq_len + + self.uniform_timestep = False + self.scheduler = FlowMatchScheduler( + shift=timestep_shift, sigma_min=0.0, extra_one_step=True + ) + self.scheduler.set_timesteps(1000, training=True) + self.seq_len = 28160 + self._compiled_model_call = None + self.get_scheduler() + + def get_scheduler(self) -> SchedulerInterface: + scheduler = self.scheduler + scheduler.convert_x0_to_noise = types.MethodType( + SchedulerInterface.convert_x0_to_noise, scheduler + ) + scheduler.convert_noise_to_x0 = types.MethodType( + SchedulerInterface.convert_noise_to_x0, scheduler + ) + scheduler.convert_velocity_to_x0 = types.MethodType( + SchedulerInterface.convert_velocity_to_x0, scheduler + ) + self.scheduler = scheduler + return scheduler + + def configure_torch_compile( + self, + *, + backend: str = "inductor", + mode: str | None = "max-autotune-no-cudagraphs", + fullgraph: bool = False, + dynamic: bool | None = False, + options: dict | None = None, + suppress_errors: bool = True, + ) -> bool: + from utils.torch_compile_utils import configure_module_call_torch_compile + + self._compiled_model_call = configure_module_call_torch_compile( + self.model, + name="SPWanDiffusionWrapper5B.model", + backend=backend, + mode=mode, + fullgraph=fullgraph, + dynamic=dynamic, + options=options, + suppress_errors=suppress_errors, + ) + return self._compiled_model_call is not None + + def _call_model(self, *args, **kwargs): + if self._compiled_model_call is not None: + return self._compiled_model_call(*args, **kwargs) + return self.model(*args, **kwargs) + + def _convert_flow_pred_to_x0( + self, flow_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor + ) -> torch.Tensor: + original_dtype = flow_pred.dtype + flow_pred, xt, sigmas, timesteps = map( + lambda x: x.double().to(flow_pred.device), + [flow_pred, xt, self.scheduler.sigmas, self.scheduler.timesteps], + ) + timestep_id = torch.argmin( + (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1 + ) + sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1) + return (xt - sigma_t * flow_pred).to(original_dtype) + + def forward( + self, + noisy_image_or_video: torch.Tensor, + conditional_dict: dict, + timestep: torch.Tensor, + kv_cache: Optional[List[dict]] = None, + crossattn_cache: Optional[List[dict]] = None, + current_start: Optional[int] = None, + cache_start: Optional[int] = None, + **_, + ) -> torch.Tensor: + prompt_embeds = conditional_dict["prompt_embeds"] + input_timestep = timestep[:, 0] if self.uniform_timestep else timestep + flow_pred = self._call_model( + noisy_image_or_video.permute(0, 2, 1, 3, 4), + t=input_timestep, + context=prompt_embeds, + seq_len=self.seq_len, + kv_cache=kv_cache, + crossattn_cache=crossattn_cache, + current_start=current_start if current_start is not None else 0, + cache_start=cache_start if cache_start is not None else 0, + ).permute(0, 2, 1, 3, 4) + pred_x0 = self._convert_flow_pred_to_x0( + flow_pred=flow_pred.flatten(0, 1), + xt=noisy_image_or_video.flatten(0, 1), + timestep=timestep.flatten(0, 1), + ).unflatten(0, flow_pred.shape[:2]) + return flow_pred, pred_x0 + + +class CausalDiffusionInferencePipelineSP(CausalDiffusionInferencePipeline): + """LongLive2.0 diffusion inference pipeline using Ulysses sequence parallelism.""" + + def __init__( + self, + args, + device, + generator=None, + text_encoder=None, + vae=None, + sp_group=None, + dp_rank: int = 0, + ): + if dist.is_initialized(): + init_sequence_parallel(group=sp_group) + self.dp_rank = dp_rank + if generator is None: + model_kwargs = getattr(args, "model_kwargs", {}) + generator = SPWanDiffusionWrapper5B( + model_name=_model_kw(model_kwargs, "model_name", "Wan2.2-TI2V-5B"), + timestep_shift=_model_kw(model_kwargs, "timestep_shift", 5.0), + local_attn_size=_model_kw(model_kwargs, "local_attn_size", -1), + sink_size=_model_kw(model_kwargs, "sink_size", 0), + num_frame_per_block=getattr(args, "num_frame_per_block", 1), + t_scale=getattr(args, "t_scale", 1.0), + rope_method=getattr(args, "rope_method", "linear"), + original_seq_len=getattr(args, "original_seq_len", None), + ) + super().__init__( + args=args, + device=device, + generator=generator, + text_encoder=text_encoder, + vae=vae, + ) + if self.quantize_kv: + raise ValueError("kv_quant is not supported in Ulysses SP inference.") + sp_print( + f"SP diffusion pipeline initialized: nfpb={self.num_frame_per_block}, " + f"sp_world={get_sp_world_size() if is_sp_enabled() else 1}, dp_rank={dp_rank}" + ) + + def _initialize_kv_cache(self, batch_size, dtype, device): + """Initialize head-split KV caches for Ulysses SP.""" + kv_cache_pos = [] + kv_cache_neg = [] + num_heads = wan_default_config[self.model_name]["num_heads"] + head_dim = wan_default_config[self.model_name]["head_dim"] + sp_world_size = get_sp_world_size() if is_sp_enabled() else 1 + if num_heads % sp_world_size != 0: + raise ValueError( + f"num_heads ({num_heads}) must be divisible by sp_world_size ({sp_world_size})" + ) + cache_heads = num_heads // sp_world_size + + if self.local_attn_size != -1: + kv_cache_size = self.local_attn_size * self.frame_seq_length + else: + kv_cache_size = 3 * self.num_frame_per_block * self.frame_seq_length + block_token_size = self.num_frame_per_block * self.frame_seq_length + max_blocks = kv_cache_size // block_token_size + + sp_print( + f"Initializing SP KV cache: size={kv_cache_size}, heads/rank={cache_heads}, " + f"block_token_size={block_token_size}" + ) + for _ in range(self.num_transformer_blocks): + entry = { + "k": torch.zeros( + [batch_size, kv_cache_size, cache_heads, head_dim], + dtype=dtype, device=device, + ), + "v": torch.zeros( + [batch_size, kv_cache_size, cache_heads, head_dim], + dtype=dtype, device=device, + ), + "quantized": False, + "block_token_size": block_token_size, + "max_blocks": max_blocks, + "num_heads": cache_heads, + "num_filled_blocks": 0, + "global_end_index": torch.tensor([0], dtype=torch.long, device=device), + "local_end_index": torch.tensor([0], dtype=torch.long, device=device), + "pinned_start": torch.tensor([-1], dtype=torch.long, device=device), + "pinned_len": torch.tensor([0], dtype=torch.long, device=device), + } + kv_cache_pos.append({key: value.clone() if torch.is_tensor(value) else value for key, value in entry.items()}) + kv_cache_neg.append({key: value.clone() if torch.is_tensor(value) else value for key, value in entry.items()}) + + self.kv_cache_pos = kv_cache_pos + self.kv_cache_neg = kv_cache_neg diff --git a/pipeline/self_forcing_training.py b/pipeline/self_forcing_training.py new file mode 100644 index 0000000..411e8ab --- /dev/null +++ b/pipeline/self_forcing_training.py @@ -0,0 +1,774 @@ +# Adopted from https://github.com/guandeh17/Self-Forcing +# SPDX-License-Identifier: Apache-2.0 +from utils.wan_5b_wrapper import WanDiffusionWrapper +from utils.scheduler import SchedulerInterface +from utils.i2v_conditioning import ( + _overwrite_i2v_context, + _zero_i2v_context_timestep, +) +from typing import List, Optional, Tuple +import torch +import torch.distributed as dist +from torchvision.io import write_video + + + +class SelfForcingTrainingPipeline: + def __init__(self, + scheduler: SchedulerInterface, + generator: WanDiffusionWrapper, + denoising_step_list: Optional[List[int]] = None, + num_frame_per_block=3, + independent_first_frame: bool = False, + same_step_across_blocks: bool = False, + last_step_only: bool = False, + num_max_frames: int = 21, + context_noise: int = 0, + sampling_steps: Optional[int] = None, + local_attn_size: int = -1, + sink_size: int = 0, + multi_shot_sink: bool = False, + scene_cut_prefix: str = "[SCENE_CUT]", + multi_shot_rope_offset: float = 0.0, + frame_seq_length: Optional[int] = None, + **kwargs): + super().__init__() + self.scheduler = scheduler + self.generator = generator + if denoising_step_list is None: + if sampling_steps is None: + raise ValueError("sampling_steps is required when denoising_step_list is not provided") + denoising_step_list = self._build_default_denoising_step_list(sampling_steps) + self.denoising_step_list = torch.as_tensor(denoising_step_list, dtype=torch.long) + if self.denoising_step_list[-1] == 0: + self.denoising_step_list = self.denoising_step_list[:-1] # remove the zero timestep for inference + + # Wan specific hyperparameters + self.num_transformer_blocks = self.generator.model.num_layers + if not dist.is_initialized() or dist.get_rank() == 0: + print(f"num_transformer_blocks: {self.num_transformer_blocks}") + if frame_seq_length is not None: + self.frame_seq_length = frame_seq_length + else: + self.frame_seq_length = 880 + if not dist.is_initialized() or dist.get_rank() == 0: + print(f"frame_seq_length: {self.frame_seq_length}") + self.num_frame_per_block = num_frame_per_block + self.context_noise = context_noise + self.i2v = False + + self.kv_cache1 = None + self.kv_cache2 = None + self.crossattn_cache = None + self.independent_first_frame = independent_first_frame + self.same_step_across_blocks = same_step_across_blocks + self.last_step_only = last_step_only + self.sampling_steps = sampling_steps + self.local_attn_size = local_attn_size + self.sink_size = sink_size + self.multi_shot_sink = multi_shot_sink + self.global_sink_size = sink_size if multi_shot_sink else 0 + self.scene_cut_prefix = scene_cut_prefix + self.multi_shot_rope_offset = multi_shot_rope_offset + self.kv_cache_size = num_max_frames * self.frame_seq_length + + if not dist.is_initialized() or dist.get_rank() == 0: + print( + f"[SelfForcingTrainingPipeline] kv_cache_size={self.kv_cache_size}, " + f"local_attn_size={local_attn_size}, sink_size={sink_size}, " + f"auto_global_sink_size={self.global_sink_size}, multi_shot_sink={multi_shot_sink}" + ) + + def _build_default_denoising_step_list(self, sampling_steps): + shift = getattr(self.scheduler, "shift", 1.0) + num_train_timesteps = getattr(self.scheduler, "num_train_timesteps", 1000) + sigmas = torch.linspace(1.0, 0.0, int(sampling_steps) + 1, dtype=torch.float32)[:-1] + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + return torch.cat([ + (sigmas * num_train_timesteps).to(torch.long), + torch.zeros(1, dtype=torch.long), + ]) + + def generate_and_sync_list(self, num_blocks, num_denoising_steps, device): + rank = dist.get_rank() if dist.is_initialized() else 0 + + if rank == 0: + # Generate random indices + indices = torch.randint( + low=0, + high=num_denoising_steps, + size=(num_blocks,), + device=device + ) + if self.last_step_only: + indices = torch.ones_like(indices) * (num_denoising_steps - 1) + else: + indices = torch.empty(num_blocks, dtype=torch.long, device=device) + if dist.is_initialized(): + dist.broadcast(indices, src=0) # Broadcast the random indices to all ranks + return indices.tolist() + + def generate_chunk_with_cache( + self, + noise: torch.Tensor, + conditional_dict: dict, + *, + current_start_frame: int = 0, + requires_grad: bool = True, + return_sim_step: bool = False, + ) -> Tuple[torch.Tensor, Optional[int], Optional[int]]: + """ + Chunk generation method tailored for sequential training + + Args: + noise: noise tensor for a single chunk [batch_size, chunk_frames, C, H, W] + conditional_dict: dictionary of conditional information + kv_cache: externally provided KV cache (defaults to self.kv_cache1 if None) + crossattn_cache: externally provided cross-attention cache (defaults to self.crossattn_cache if None) + current_start_frame: start frame index of the chunk in the full sequence + requires_grad: whether gradients are required + return_sim_step: whether to return simulation step info + + Returns: + output: generated chunk [batch_size, chunk_frames, C, H, W] + denoised_timestep_from: starting denoise timestep + denoised_timestep_to: ending denoise timestep + """ + batch_size, chunk_frames, num_channels, height, width = noise.shape + + # Compute block configuration + if not self.independent_first_frame or chunk_frames % self.num_frame_per_block == 0: + assert chunk_frames % self.num_frame_per_block == 0 + num_blocks = chunk_frames // self.num_frame_per_block + all_num_frames = [self.num_frame_per_block] * num_blocks + else: + # Handle the case of an independent first frame + assert (chunk_frames - 1) % self.num_frame_per_block == 0 + num_blocks = (chunk_frames - 1) // self.num_frame_per_block + all_num_frames = [1] + [self.num_frame_per_block] * num_blocks + + # Prepare output tensor + output = torch.zeros_like(noise) + + # Build per-block conditional dicts for prompt switching + prompt_embeds = conditional_dict["prompt_embeds"] + num_prompts = prompt_embeds.shape[0] + num_segments = num_prompts // batch_size + if num_segments > 1: + prompt_embeds_per_block = prompt_embeds.reshape( + batch_size, num_segments, *prompt_embeds.shape[1:]) + conditional_dict_list = [ + {"prompt_embeds": prompt_embeds_per_block[:, i]} + for i in range(num_segments) + ] + else: + conditional_dict_list = None + + # Randomly select denoising steps (synced across ranks) + num_denoising_steps = len(self.denoising_step_list) + exit_flags = self.generate_and_sync_list(len(all_num_frames), num_denoising_steps, device=noise.device) + + # Determine gradient-enabled range — disable everywhere when requires_grad=False + if not requires_grad: + start_gradient_frame_index = chunk_frames # Out of range: no gradients anywhere + else: + start_gradient_frame_index = 0 + + if not dist.is_initialized() or dist.get_rank() == 0: + print(f"[SeqTrain-Pipeline] start_gradient_frame_index={start_gradient_frame_index}, chunk_frames={chunk_frames}") + + # Generate block by block + local_start_frame = 0 + for block_index, current_num_frames in enumerate(all_num_frames): + if conditional_dict_list is not None: + block_cond = conditional_dict_list[min(block_index, len(conditional_dict_list) - 1)] + for cache_idx in range(self.num_transformer_blocks): + self.crossattn_cache[cache_idx]["is_init"] = False + else: + block_cond = conditional_dict + + noisy_input = noise[:, local_start_frame:local_start_frame + current_num_frames] + + # Spatial denoising loop + for step_idx, current_timestep in enumerate(self.denoising_step_list): + exit_flag = ( + step_idx == exit_flags[0] + if self.same_step_across_blocks + else step_idx == exit_flags[block_index] + ) + + timestep = torch.ones( + [batch_size, current_num_frames], + device=noise.device, + dtype=torch.int64 + ) * current_timestep + + if not exit_flag: + # Intermediate steps: no gradients + with torch.no_grad(): + _, denoised_pred = self.generator( + noisy_image_or_video=noisy_input, + conditional_dict=block_cond, + timestep=timestep, + kv_cache=self.kv_cache1, + crossattn_cache=self.crossattn_cache, + current_start=(current_start_frame + local_start_frame) * self.frame_seq_length, + ) + + # Add noise for the next step + if step_idx < len(self.denoising_step_list) - 1: + next_timestep = self.denoising_step_list[step_idx + 1] + noisy_input = self.scheduler.add_noise( + denoised_pred.flatten(0, 1), + torch.randn_like(denoised_pred.flatten(0, 1)), + next_timestep * torch.ones( + [batch_size * current_num_frames], device=noise.device, dtype=torch.long + ), + ).unflatten(0, denoised_pred.shape[:2]) + else: + # Final step may require gradients + enable_grad = local_start_frame >= start_gradient_frame_index + + context_manager = torch.enable_grad() if enable_grad else torch.no_grad() + with context_manager: + _, denoised_pred = self.generator( + noisy_image_or_video=noisy_input, + conditional_dict=block_cond, + timestep=timestep, + kv_cache=self.kv_cache1, + crossattn_cache=self.crossattn_cache, + current_start=(current_start_frame + local_start_frame) * self.frame_seq_length, + ) + break + + # Record output + output[:, local_start_frame:local_start_frame + current_num_frames] = denoised_pred + + # Update cache with context noise + context_timestep = torch.ones_like(timestep) * self.context_noise + context_noisy = self.scheduler.add_noise( + denoised_pred.flatten(0, 1), + torch.randn_like(denoised_pred.flatten(0, 1)), + context_timestep.flatten(0, 1), + ).unflatten(0, denoised_pred.shape[:2]) + + with torch.no_grad(): + self.generator( + noisy_image_or_video=context_noisy, + conditional_dict=block_cond, + timestep=context_timestep, + kv_cache=self.kv_cache1, + crossattn_cache=self.crossattn_cache, + current_start=(current_start_frame + local_start_frame) * self.frame_seq_length, + ) + + local_start_frame += current_num_frames + + # Compute returned timestep information + if not self.same_step_across_blocks: + denoised_timestep_from, denoised_timestep_to = None, None + elif exit_flags[0] == len(self.denoising_step_list) - 1: + denoised_timestep_to = 0 + denoised_timestep_from = 1000 - torch.argmin( + (self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0]].cuda()).abs(), dim=0 + ).item() + else: + denoised_timestep_to = 1000 - torch.argmin( + (self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0] + 1].cuda()).abs(), dim=0 + ).item() + denoised_timestep_from = 1000 - torch.argmin( + (self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0]].cuda()).abs(), dim=0 + ).item() + + if return_sim_step: + return output, denoised_timestep_from, denoised_timestep_to, exit_flags[0] + 1 + + return output, denoised_timestep_from, denoised_timestep_to + + def inference_with_trajectory( + self, + noise: torch.Tensor, + initial_latent: Optional[torch.Tensor] = None, + return_sim_step: bool = False, + slice_last_frames: int = 21, + sampling_steps: Optional[int] = None, + **conditional_dict + ) -> torch.Tensor: + # Apply local_attn_size / sink_size overrides before inference, + # matching what CausalDiffusionInferencePipeline does. + prev_state = self._apply_attn_overrides() + try: + return self._inference_with_trajectory_inner( + noise=noise, + initial_latent=initial_latent, + return_sim_step=return_sim_step, + slice_last_frames=slice_last_frames, + sampling_steps=sampling_steps, + **conditional_dict, + ) + finally: + self._restore_attn_overrides(prev_state) + + def _apply_attn_overrides(self): + """Save current model attention state and apply pipeline overrides.""" + model = self.generator.model + prev = { + "local_attn_size": getattr(model, "local_attn_size", -1), + "rope_temporal_offset": getattr(model, "rope_temporal_offset", 0.0), + "max_attention_sizes": {}, + "sink_sizes": {}, + "global_sink_sizes": {}, + } + for name, module in model.named_modules(): + if hasattr(module, "max_attention_size"): + prev["max_attention_sizes"][name] = module.max_attention_size + if hasattr(module, "sink_size"): + prev["sink_sizes"][name] = module.sink_size + if hasattr(module, "global_sink_size"): + prev["global_sink_sizes"][name] = module.global_sink_size + + model.local_attn_size = self.local_attn_size + model.rope_temporal_offset = 0.0 + self._set_all_modules_max_attention_size(self.local_attn_size) + if self.sink_size is not None and self.sink_size >= 0: + self._set_all_modules_sink_size(self.sink_size) + self._set_all_modules_global_sink_size(self.global_sink_size) + + return prev + + def _restore_attn_overrides(self, prev): + """Restore model attention state saved by _apply_attn_overrides.""" + model = self.generator.model + model.local_attn_size = prev["local_attn_size"] + model.rope_temporal_offset = prev["rope_temporal_offset"] + for name, module in model.named_modules(): + if name in prev["max_attention_sizes"]: + try: + module.max_attention_size = prev["max_attention_sizes"][name] + except Exception: + pass + if name in prev["sink_sizes"]: + try: + module.sink_size = prev["sink_sizes"][name] + except Exception: + pass + if name in prev["global_sink_sizes"]: + try: + module.global_sink_size = prev["global_sink_sizes"][name] + except Exception: + pass + + @staticmethod + def _is_scene_cut_from_mask(scene_cut_mask, block_index: int) -> bool: + if scene_cut_mask is None or block_index <= 0: + return False + if block_index >= len(scene_cut_mask): + return False + value = scene_cut_mask[block_index] + if torch.is_tensor(value): + return bool(value.item()) + return bool(value) + + def _set_all_modules_sink_size(self, sink_size_value: int): + """Override sink_size on all submodules that define it.""" + model = self.generator.model + if hasattr(model, "sink_size"): + model.sink_size = sink_size_value + for _name, module in model.named_modules(): + if hasattr(module, "sink_size"): + try: + module.sink_size = sink_size_value + except Exception: + pass + + def _set_all_modules_global_sink_size(self, value: int): + """Override global_sink_size on all submodules; create the attribute if missing.""" + setattr(self.generator.model, "global_sink_size", value) + for _, module in self.generator.model.named_modules(): + try: + setattr(module, "global_sink_size", value) + except Exception: + pass + + def _pin_current_chunk(self, kv_cache, current_num_frames): + """Mark the current chunk's buffer position as pinned for multi-shot sink. + + The pinned region REPLACES the original sink on the next rolling event. + No data is copied here — relocation happens inside the attention layer + during rolling, ensuring zero duplication. + """ + chunk_tokens = current_num_frames * self.frame_seq_length + pin_len = min(self.sink_size * self.frame_seq_length, chunk_tokens) + + for block_cache in kv_cache: + local_end = block_cache["local_end_index"].item() + chunk_start = local_end - chunk_tokens + block_cache["pinned_start"].fill_(chunk_start) + block_cache["pinned_len"].fill_(pin_len) + + def _inference_with_trajectory_inner( + self, + noise: torch.Tensor, + initial_latent: Optional[torch.Tensor] = None, + return_sim_step: bool = False, + slice_last_frames: int = 21, + sampling_steps: Optional[int] = None, + **conditional_dict + ) -> torch.Tensor: + from wan_5b.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler + + batch_size, num_frames, num_channels, height, width = noise.shape + num_input_frames = initial_latent.shape[1] if initial_latent is not None else 0 + clamp_i2v_first_chunk = self.independent_first_frame and initial_latent is not None + if clamp_i2v_first_chunk and num_input_frames != 1: + raise ValueError( + f"i2v first-chunk clamp expects one conditioning latent frame, got {num_input_frames}." + ) + + if not self.independent_first_frame or clamp_i2v_first_chunk: + # If the first frame is independent and the first frame is provided, then the number of frames in the + # noise should still be a multiple of num_frame_per_block + assert num_frames % self.num_frame_per_block == 0 + num_blocks = num_frames // self.num_frame_per_block + else: + # Using a [1, 4, 4, 4, 4, 4, ...] model to generate a video without image conditioning + assert (num_frames - 1) % self.num_frame_per_block == 0 + num_blocks = (num_frames - 1) // self.num_frame_per_block + num_output_frames = ( + num_frames if clamp_i2v_first_chunk else num_frames + num_input_frames + ) + output = torch.zeros( + [batch_size, num_output_frames, num_channels, height, width], + device=noise.device, + dtype=noise.dtype + ) + + # Step 1: Initialize KV cache to all zeros + self._initialize_kv_cache( + batch_size=batch_size, dtype=noise.dtype, device=noise.device + ) + self._initialize_crossattn_cache( + batch_size=batch_size, dtype=noise.dtype, device=noise.device + ) + + # Build per-block conditional dicts for prompt switching + prompt_embeds = conditional_dict["prompt_embeds"] + num_prompts = prompt_embeds.shape[0] + num_segments = num_prompts // batch_size + if num_segments > 1: + prompt_embeds_per_block = prompt_embeds.reshape( + batch_size, num_segments, *prompt_embeds.shape[1:]) + conditional_dict_list = [ + {"prompt_embeds": prompt_embeds_per_block[:, i]} + for i in range(num_segments) + ] + else: + conditional_dict_list = None + + # Step 2: Cache context feature + current_start_frame = 0 + if initial_latent is not None and not clamp_i2v_first_chunk: + timestep = torch.ones([batch_size, 1], device=noise.device, dtype=torch.int64) * 0 + # Assume num_input_frames is 1 + self.num_frame_per_block * num_input_blocks + output[:, :1] = initial_latent + init_cond = conditional_dict_list[0] if conditional_dict_list is not None else conditional_dict + with torch.no_grad(): + self.generator( + noisy_image_or_video=initial_latent, + conditional_dict=init_cond, + timestep=timestep * 0, + kv_cache=self.kv_cache1, + crossattn_cache=self.crossattn_cache, + current_start=current_start_frame * self.frame_seq_length + ) + current_start_frame += 1 + + # Step 3: Temporal denoising loop + all_num_frames = [self.num_frame_per_block] * num_blocks + if self.independent_first_frame and initial_latent is None: + all_num_frames = [1] + all_num_frames + + # --- UniPC scheduler setup --- + # Priority: function arg > pipeline attribute > len(denoising_step_list) + if sampling_steps is None: + sampling_steps = self.sampling_steps if self.sampling_steps is not None else len(self.denoising_step_list) + shift = self.scheduler.shift + num_train_timesteps = self.scheduler.num_train_timesteps + ref_scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=num_train_timesteps, shift=1, use_dynamic_shifting=False) + ref_scheduler.set_timesteps(sampling_steps, device=noise.device, shift=shift) + unipc_timesteps = ref_scheduler.timesteps + num_denoising_steps = len(unipc_timesteps) + + exit_flags = self.generate_and_sync_list(len(all_num_frames), num_denoising_steps, device=noise.device) + if slice_last_frames == -1: + # -1 means keep full sequence trainable. + start_gradient_frame_index = 0 + else: + start_gradient_frame_index = num_output_frames - slice_last_frames + + scene_cut_mask = conditional_dict.pop("scene_cut_mask", None) + current_shot_index = 0 + phi = self.multi_shot_rope_offset + self.generator.model.rope_temporal_offset = 0.0 + + grad_enable_mask = torch.zeros((batch_size, sum(all_num_frames)), dtype=torch.bool) + for block_index, current_num_frames in enumerate(all_num_frames): + if phi != 0.0 and self._is_scene_cut_from_mask(scene_cut_mask, block_index): + current_shot_index += 1 + self.generator.model.rope_temporal_offset = current_shot_index * phi + if not dist.is_initialized() or dist.get_rank() == 0: + print( + f"[training] multi-shot RoPE: shot_index={current_shot_index}, " + f"temporal_offset={self.generator.model.rope_temporal_offset:.4f}" + ) + + if conditional_dict_list is not None: + block_cond = conditional_dict_list[min(block_index, len(conditional_dict_list) - 1)] + for cache_idx in range(self.num_transformer_blocks): + self.crossattn_cache[cache_idx]["is_init"] = False + else: + block_cond = conditional_dict + + first_i2v_block = clamp_i2v_first_chunk and block_index == 0 + noise_start_frame = ( + current_start_frame + if clamp_i2v_first_chunk + else current_start_frame - num_input_frames + ) + latents = noise[ + :, + noise_start_frame:noise_start_frame + current_num_frames, + ] + + # re-init scheduler per chunk (internal state is consumed during stepping) + sample_scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=num_train_timesteps, shift=1, use_dynamic_shifting=False) + sample_scheduler.set_timesteps(sampling_steps, device=noise.device, shift=shift) + + # Step 3.1: Spatial denoising loop (UniPC multi-step) + for index, t in enumerate(sample_scheduler.timesteps): + if self.same_step_across_blocks: + exit_flag = (index == exit_flags[0]) + else: + exit_flag = (index == exit_flags[block_index]) + timestep = t * torch.ones( + [batch_size, current_num_frames], + device=noise.device, + dtype=torch.float32) + if first_i2v_block: + latents = _overwrite_i2v_context( + latents, initial_latent, num_input_frames + ) + timestep = _zero_i2v_context_timestep( + timestep, num_input_frames + ) + if not exit_flag: + with torch.no_grad(): + flow_pred, _ = self.generator( + noisy_image_or_video=latents, + conditional_dict=block_cond, + timestep=timestep, + kv_cache=self.kv_cache1, + crossattn_cache=self.crossattn_cache, + current_start=current_start_frame * self.frame_seq_length + ) + latents = sample_scheduler.step( + flow_pred, t, latents, return_dict=False)[0] + if first_i2v_block: + latents = _overwrite_i2v_context( + latents, initial_latent, num_input_frames + ) + else: + if current_start_frame < start_gradient_frame_index: + grad_enable_mask[:, current_start_frame:current_start_frame + current_num_frames] = False + with torch.no_grad(): + flow_pred, denoised_pred = self.generator( + noisy_image_or_video=latents, + conditional_dict=block_cond, + timestep=timestep, + kv_cache=self.kv_cache1, + crossattn_cache=self.crossattn_cache, + current_start=current_start_frame * self.frame_seq_length + ) + else: + grad_enable_mask[:, current_start_frame:current_start_frame + current_num_frames] = True + flow_pred, denoised_pred = self.generator( + noisy_image_or_video=latents, + conditional_dict=block_cond, + timestep=timestep, + kv_cache=self.kv_cache1, + crossattn_cache=self.crossattn_cache, + current_start=current_start_frame * self.frame_seq_length + ) + if first_i2v_block: + denoised_pred = _overwrite_i2v_context( + denoised_pred, initial_latent, num_input_frames + ) + break + + # Step 3.2: record the model's output + output[:, current_start_frame:current_start_frame + current_num_frames] = denoised_pred + + # Step 3.3: rerun with context noise to update the cache + context_timestep = torch.ones( + [batch_size, current_num_frames], device=noise.device, dtype=torch.long) * self.context_noise + if first_i2v_block: + context_timestep = _zero_i2v_context_timestep( + context_timestep, num_input_frames + ) + # add context noise + context_noise = torch.randn_like(denoised_pred.flatten(0, 1)) + if first_i2v_block: + context_noise = context_noise.unflatten(0, denoised_pred.shape[:2]) + context_noise[:, :num_input_frames] = 0 + context_noise = context_noise.flatten(0, 1) + denoised_pred = self.scheduler.add_noise( + denoised_pred.flatten(0, 1), + context_noise, + context_timestep.reshape(1, -1) * torch.ones( + [batch_size * current_num_frames], device=noise.device, dtype=torch.long) + ).unflatten(0, denoised_pred.shape[:2]) + if first_i2v_block: + denoised_pred = _overwrite_i2v_context( + denoised_pred, initial_latent, num_input_frames + ) + with torch.no_grad(): + self.generator( + noisy_image_or_video=denoised_pred, + conditional_dict=block_cond, + timestep=context_timestep, + kv_cache=self.kv_cache1, + crossattn_cache=self.crossattn_cache, + current_start=current_start_frame * self.frame_seq_length + ) + + # Step 3.3b: pin KV on scene cut for multi-shot sink. + if self.multi_shot_sink and scene_cut_mask is not None: + is_cut = ( + block_index > 0 + and block_index < len(scene_cut_mask) + and scene_cut_mask[block_index] + ) + if is_cut: + self._pin_current_chunk(self.kv_cache1, current_num_frames) + + # Step 3.4: update the start and end frame indices + current_start_frame += current_num_frames + + # Step 3.5: Return the denoised timestep + if not self.same_step_across_blocks: + denoised_timestep_from, denoised_timestep_to = None, None + elif exit_flags[0] == num_denoising_steps - 1: + denoised_timestep_to = 0 + denoised_timestep_from = 1000 - torch.argmin( + (self.scheduler.timesteps.cuda() - unipc_timesteps[exit_flags[0]].cuda()).abs(), dim=0).item() + else: + denoised_timestep_to = 1000 - torch.argmin( + (self.scheduler.timesteps.cuda() - unipc_timesteps[exit_flags[0] + 1].cuda()).abs(), dim=0).item() + denoised_timestep_from = 1000 - torch.argmin( + (self.scheduler.timesteps.cuda() - unipc_timesteps[exit_flags[0]].cuda()).abs(), dim=0).item() + + if return_sim_step: + return output, denoised_timestep_from, denoised_timestep_to, exit_flags[0] + 1 + + return output, denoised_timestep_from, denoised_timestep_to + + def _initialize_kv_cache(self, batch_size, dtype, device): + """ + Initialize a Per-GPU KV cache for the Wan model. + """ + kv_cache1 = [] + # Get the actual number of heads and head dimension from model + num_heads = self.generator.model.num_heads + head_dim = self.generator.model.dim // num_heads + + for _ in range(self.num_transformer_blocks): + kv_cache1.append({ + "k": torch.zeros([batch_size, self.kv_cache_size, num_heads, head_dim], dtype=dtype, device=device), + "v": torch.zeros([batch_size, self.kv_cache_size, num_heads, head_dim], dtype=dtype, device=device), + "global_end_index": torch.tensor([0], dtype=torch.long, device=device), + "local_end_index": torch.tensor([0], dtype=torch.long, device=device), + "pinned_start": torch.tensor([0], dtype=torch.long, device=device), + "pinned_len": torch.tensor([0], dtype=torch.long, device=device), + }) + + self.kv_cache1 = kv_cache1 + + def _initialize_crossattn_cache(self, batch_size, dtype, device): + """ + Initialize a Per-GPU cross-attention cache for the Wan model. + """ + crossattn_cache = [] + + # Get the actual number of heads and head dimension from model + num_heads = self.generator.model.num_heads + head_dim = self.generator.model.dim // num_heads + + for _ in range(self.num_transformer_blocks): + crossattn_cache.append({ + "k": torch.zeros([batch_size, 512, num_heads, head_dim], dtype=dtype, device=device), + "v": torch.zeros([batch_size, 512, num_heads, head_dim], dtype=dtype, device=device), + "is_init": False + }) + + self.crossattn_cache = crossattn_cache + + def clear_kv_cache(self): + """ + Zero out all tensors in KV cache and cross-attention cache instead of setting them to None. + This preserves memory allocation while clearing old information, avoiding reallocation overhead. + """ + + # Clear KV cache + if getattr(self, "kv_cache1", None) is not None: + for blk in self.kv_cache1: + blk["k"].zero_() + blk["v"].zero_() + if "global_end_index" in blk: + blk["global_end_index"].zero_() + if "local_end_index" in blk: + blk["local_end_index"].zero_() + if "pinned_start" in blk: + blk["pinned_start"].zero_() + if "pinned_len" in blk: + blk["pinned_len"].zero_() + + # Clear cross-attention cache + if getattr(self, "crossattn_cache", None) is not None: + for blk in self.crossattn_cache: + blk["k"].zero_() + blk["v"].zero_() + blk["is_init"] = False + + def _set_all_modules_max_attention_size(self, local_attn_size_value: int): + """ + Set a unified upper bound for all submodules that contain the max_attention_size attribute. + local_attn_size_value == -1 indicates global attention (use Wan's default token limit 32760). + Otherwise set to local_attn_size_value * frame_seq_length. + """ + if isinstance(local_attn_size_value, (list, tuple)): + raise ValueError("_set_all_modules_max_attention_size expects an int, got list/tuple.") + + if int(local_attn_size_value) == -1: + target_size = 32760 + policy = "global" + else: + target_size = int(local_attn_size_value) * self.frame_seq_length + policy = "local" + + # Root module + if hasattr(self.generator.model, "max_attention_size"): + try: + _ = getattr(self.generator.model, "max_attention_size") + except Exception: + pass + setattr(self.generator.model, "max_attention_size", target_size) + + # Child modules + for name, module in self.generator.model.named_modules(): + if hasattr(module, "max_attention_size"): + try: + setattr(module, "max_attention_size", target_size) + except Exception: + pass diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..315a68e --- /dev/null +++ b/requirements.txt @@ -0,0 +1,42 @@ +opencv-python>=4.9.0.80 +diffusers==0.31.0 +transformers>=4.49.0,<5 +tokenizers>=0.20.3 +accelerate>=1.1.1 +tqdm +datasets +imageio +easydict +ftfy +dashscope +imageio-ffmpeg +# numpy==1.24.4 +wandb +omegaconf +einops +av==13.1.0 +opencv-python +git+https://github.com/openai/CLIP.git +open_clip_torch +starlette +pycocotools +lmdb +matplotlib +sentencepiece +pydantic==2.10.6 +scikit-image +huggingface_hub[cli] +dominate +# Optional TensorRT support. Install separately because nvidia-tensorrt is a +# placeholder on PyPI and requires NVIDIA's package index to be configured first: +# pip install nvidia-pyindex +# pip install nvidia-tensorrt +# pip install pycuda +onnx +onnxruntime +onnxscript +onnxconverter_common +flask +flask-socketio +torchao==0.13.0 # paired with the documented PyTorch 2.8.0 environment +peft diff --git a/scripts/compute_sp_vae_chunk_halo.py b/scripts/compute_sp_vae_chunk_halo.py new file mode 100644 index 0000000..e6ded23 --- /dev/null +++ b/scripts/compute_sp_vae_chunk_halo.py @@ -0,0 +1,163 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0 +# +# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Compute SP chunk-halo VAE frame windows for a training config.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +from omegaconf import OmegaConf + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from utils.config import normalize_config, wan_default_config + +DEFAULT_SP_VAE_HALO_LATENTS = 28 + + +def latent_range_to_raw_window( + latent_start: int, + latent_end: int, + *, + temporal_compression_ratio: int, +) -> tuple[int, int, int]: + """Map a latent range to the raw-frame window needed by Wan VAE encode.""" + if latent_end <= latent_start: + raise ValueError(f"latent_end must be > latent_start, got {latent_start}, {latent_end}") + ratio = int(temporal_compression_ratio) + if latent_start == 0: + return 0, 1 + ratio * (latent_end - 1), 0 + return ratio * (latent_start - 1), 1 + ratio * (latent_end - 1), 1 + + +def compute_chunk_halo_metas( + *, + total_latent_frames: int, + total_raw_frames: int, + sp_size: int, + halo_latents: int, + temporal_compression_ratio: int, +) -> list[dict[str, int]]: + if sp_size <= 0: + raise ValueError("sp_size must be positive") + if halo_latents < 0: + raise ValueError("halo_latents must be non-negative") + if total_latent_frames % sp_size != 0: + raise ValueError( + f"total_latent_frames={total_latent_frames} must be divisible by sp_size={sp_size}" + ) + + local_latent_frames = total_latent_frames // sp_size + metas = [] + for sp_rank in range(sp_size): + keep_start = sp_rank * local_latent_frames + keep_end = keep_start + local_latent_frames + halo_start = max(0, keep_start - halo_latents) + raw_start, raw_end, pseudo_prefix_latents = latent_range_to_raw_window( + halo_start, + keep_end, + temporal_compression_ratio=temporal_compression_ratio, + ) + raw_start = max(0, raw_start) + raw_end = min(total_raw_frames, raw_end) + drop_latents = pseudo_prefix_latents + (keep_start - halo_start) + metas.append( + { + "sp_rank": sp_rank, + "keep_start": keep_start, + "keep_end": keep_end, + "halo_start": halo_start, + "raw_start": raw_start, + "raw_end": raw_end, + "raw_frames": raw_end - raw_start, + "drop_latents": drop_latents, + "local_latent_frames": local_latent_frames, + } + ) + return metas + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", default="configs/train_ar.yaml", help="Training config path") + parser.add_argument("--sp-size", type=int, default=None, help="Override sequence_parallel_size") + parser.add_argument("--halo-latents", type=int, default=None, help="Override vae_halo_latents") + parser.add_argument("--latent-frames", type=int, default=None, help="Override image_or_video_shape[1]") + parser.add_argument("--raw-frames", type=int, default=None, help="Override total raw frame count") + parser.add_argument("--json", action="store_true", help="Print JSON instead of a table") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + cfg = normalize_config(OmegaConf.load(args.config)) + model_name = cfg.model_kwargs.model_name + temporal_ratio = int(wan_default_config[model_name]["temporal_compression_ratio"]) + total_latent_frames = int( + args.latent_frames if args.latent_frames is not None else cfg.image_or_video_shape[1] + ) + total_raw_frames = int( + args.raw_frames + if args.raw_frames is not None + else ((total_latent_frames - 1) * temporal_ratio + 1) + ) + sp_size = int( + args.sp_size if args.sp_size is not None else getattr(cfg, "sequence_parallel_size", 1) + ) + halo_latents = int( + args.halo_latents + if args.halo_latents is not None + else getattr(cfg, "vae_halo_latents", DEFAULT_SP_VAE_HALO_LATENTS) + ) + + metas = compute_chunk_halo_metas( + total_latent_frames=total_latent_frames, + total_raw_frames=total_raw_frames, + sp_size=sp_size, + halo_latents=halo_latents, + temporal_compression_ratio=temporal_ratio, + ) + payload = { + "config": str(Path(args.config)), + "model_name": model_name, + "temporal_compression_ratio": temporal_ratio, + "total_latent_frames": total_latent_frames, + "total_raw_frames": total_raw_frames, + "sp_size": sp_size, + "halo_latents": halo_latents, + "rank_metas": metas, + } + if args.json: + print(json.dumps(payload, indent=2)) + return + + print( + f"config={payload['config']} model={model_name} " + f"latent_frames={total_latent_frames} raw_frames={total_raw_frames} " + f"sp_size={sp_size} halo_latents={halo_latents}" + ) + print("rank keep_latents halo_start raw_window raw_frames drop_latents") + for meta in metas: + keep_latents = f"[{meta['keep_start']},{meta['keep_end']})".ljust(12) + raw_window = f"[{meta['raw_start']},{meta['raw_end']})".ljust(11) + print( + f"{meta['sp_rank']:>4} {keep_latents} {meta['halo_start']:>10} " + f"{raw_window} {meta['raw_frames']:>10} {meta['drop_latents']:>12}" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/decode_lightvae_latents.py b/scripts/decode_lightvae_latents.py new file mode 100644 index 0000000..d2570e1 --- /dev/null +++ b/scripts/decode_lightvae_latents.py @@ -0,0 +1,346 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0 +# +# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied. +# +# SPDX-License-Identifier: Apache-2.0 + +""" +Decode saved latent files (.pt) to pixel video using a local LightVAE loader. + +This mirrors `scripts/decode_vae_latents.py`, but only depends on the current +repo's 5B VAE code plus a LightVAE checkpoint such as `MG-LightVAE_v2.pth`. + +Examples: + python scripts/decode_lightvae_latents.py \ + --input_dir /path/to/latents \ + --vae_path /path/to/MG-LightVAE_v2.pth + + torchrun --nproc_per_node=8 scripts/decode_lightvae_latents.py \ + --input_dir /path/to/latents \ + --ckpt_dir /path/to/lightvae_ckpts \ + --vae_type mg_lightvae +""" + +import argparse +import glob +import os +import sys +from typing import Optional + +os.environ.setdefault("OPENBLAS_NUM_THREADS", "1") + +import torch +import torch.distributed as dist +from torchvision.io import write_video +from tqdm import tqdm + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + + +from utils.lightvae_5b_wrapper import LightVAE5BWrapper + + +def init_distributed(): + """Initialize distributed process group if launched via torchrun.""" + if "RANK" in os.environ and "WORLD_SIZE" in os.environ: + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend="nccl") + return rank, world_size, local_rank + return 0, 1, 0 + + +def decode_latent_to_video( + vae, latent: torch.Tensor, device: torch.device, dtype: torch.dtype +) -> torch.Tensor: + """ + Decode latent to pixel video. + + Args: + vae: VAE wrapper exposing `decode_to_pixel(latent)` + latent: shape (batch, T, C, H, W) or (T, C, H, W) + device, dtype: device and dtype for computation + + Returns: + video: shape (batch, T, C, H, W), range [0, 1] + """ + if latent.dim() == 4: + latent = latent.unsqueeze(0) + latent = latent.to(device=device, dtype=dtype) + video = vae.decode_to_pixel(latent) + video = (video * 0.5 + 0.5).clamp(0, 1) + return video + + +def _normalize_requested_vae_type(value: str) -> str: + normalized = value.strip().lower() + if normalized == "wan": + return "wan2.2" + if normalized in {"wan2.2", "mg_lightvae", "mg_lightvae_v2"}: + return normalized + raise ValueError( + f"Unsupported --vae_type '{value}'. " + "Expected one of: wan2.2, wan, mg_lightvae, mg_lightvae_v2." + ) + + +def _parse_lightvae_pruning_rate(value: Optional[str]) -> Optional[float]: + if value is None: + return None + lowered = str(value).strip().lower() + if lowered in {"", "auto", "none"}: + return None + return float(lowered) + + +def _resolve_vae_paths( + *, + ckpt_dir: Optional[str], + vae_path: Optional[str], + requested_vae_type: str, + lightvae_pruning_rate: Optional[str], +): + pruning_rate = _parse_lightvae_pruning_rate(lightvae_pruning_rate) + ckpt_dir = os.path.abspath(ckpt_dir) if ckpt_dir else None + + if vae_path is not None: + resolved_vae_path = os.path.abspath(vae_path) + if requested_vae_type == "wan2.2" and pruning_rate is None: + pruning_rate = 0.0 + else: + if ckpt_dir is None: + raise ValueError("Either --ckpt_dir or --vae_path must be provided.") + if requested_vae_type == "mg_lightvae": + resolved_vae_path = os.path.join(ckpt_dir, "MG-LightVAE.pth") + if pruning_rate is None: + pruning_rate = 0.5 + elif requested_vae_type == "mg_lightvae_v2": + resolved_vae_path = os.path.join(ckpt_dir, "MG-LightVAE_v2.pth") + if pruning_rate is None: + pruning_rate = 0.75 + else: + resolved_vae_path = os.path.join(ckpt_dir, "Wan2.2_VAE.pth") + if pruning_rate is None: + pruning_rate = 0.0 + + return resolved_vae_path, pruning_rate + + +def _load_latent_tensor(pt_path: str): + try: + data = torch.load(pt_path, map_location="cpu", weights_only=True) + except TypeError: + data = torch.load(pt_path, map_location="cpu") + + if isinstance(data, torch.Tensor): + return data + if isinstance(data, dict): + latent = next(iter(data.values())) + if isinstance(latent, torch.Tensor): + return latent + raise TypeError(f"First dict value is not a tensor: {type(latent)}") + raise TypeError(f"Unsupported latent file payload type: {type(data)}") + + +def _parse_dtype(dtype_name: str) -> torch.dtype: + mapping = { + "float32": torch.float32, + "float16": torch.float16, + "bfloat16": torch.bfloat16, + } + try: + return mapping[dtype_name] + except KeyError as exc: + raise ValueError( + f"Unsupported --dtype '{dtype_name}'. " + "Expected one of: float32, float16, bfloat16." + ) from exc + + +def main(): + parser = argparse.ArgumentParser( + description="Decode saved latents to video with a local LightVAE loader." + ) + parser.add_argument( + "--input_dir", + type=str, + required=True, + help="Directory containing .pt latent files.", + ) + parser.add_argument( + "--output_dir", + type=str, + default=None, + help="Directory to save decoded .mp4 videos. Default: input_dir/decoded_lightvae_videos", + ) + parser.add_argument( + "--fps", + type=int, + default=24, + help="FPS for output video (default: 24).", + ) + parser.add_argument( + "--pattern", + type=str, + default="latents_*.pt", + help="Glob pattern for latent files (default: latents_*.pt).", + ) + parser.add_argument( + "--ckpt_dir", + type=str, + default=None, + help="Directory containing `MG-LightVAE*.pth` or `Wan2.2_VAE.pth`.", + ) + parser.add_argument( + "--vae_path", + type=str, + default=None, + help=( + "Explicit LightVAE/Wan2.2 checkpoint path. " + "Use this when you only have a single VAE checkpoint file." + ), + ) + parser.add_argument( + "--matrix_game_root", + type=str, + default=None, + help=argparse.SUPPRESS, + ) + parser.add_argument( + "--teacher_vae_path", + "--lightvae_encoder_path", + dest="teacher_vae_path", + type=str, + default=None, + help=argparse.SUPPRESS, + ) + parser.add_argument( + "--vae_type", + type=str, + default="mg_lightvae_v2", + choices=["wan2.2", "wan", "mg_lightvae", "mg_lightvae_v2"], + help=( + "VAE variant to use. " + "`mg_lightvae` maps to MG-LightVAE.pth, " + "`mg_lightvae_v2` maps to MG-LightVAE_v2.pth." + ), + ) + parser.add_argument( + "--lightvae_pruning_rate", + type=str, + default=None, + help=( + "Override pruning rate. Use `auto`/`none` to let Wan2_2_VAE infer it " + "when an explicit LightVAE checkpoint is provided." + ), + ) + parser.add_argument( + "--dtype", + type=str, + default="bfloat16", + choices=["float32", "float16", "bfloat16"], + help="Decode dtype (default: bfloat16).", + ) + args = parser.parse_args() + + rank, world_size, local_rank = init_distributed() + is_main = rank == 0 + + output_dir = args.output_dir or os.path.join( + args.input_dir, "decoded_lightvae_videos" + ) + if is_main: + os.makedirs(output_dir, exist_ok=True) + if world_size > 1: + dist.barrier() + + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + torch.set_grad_enabled(False) + + device = torch.device( + f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu" + ) + dtype = _parse_dtype(args.dtype) + + requested_vae_type = _normalize_requested_vae_type(args.vae_type) + resolved_vae_path, resolved_pruning_rate = _resolve_vae_paths( + ckpt_dir=args.ckpt_dir, + vae_path=args.vae_path, + requested_vae_type=requested_vae_type, + lightvae_pruning_rate=args.lightvae_pruning_rate, + ) + + vae = LightVAE5BWrapper( + vae_path=resolved_vae_path, + pruning_rate=resolved_pruning_rate, + device=device, + dtype=dtype, + ).eval() + + if is_main: + print( + "Using local VAE-only loader: " + f"requested={requested_vae_type}, " + f"pruning={vae.pruning_rate}, " + f"vae_path={vae.vae_path}", + flush=True, + ) + + search_path = os.path.join(args.input_dir, args.pattern) + pt_files = sorted(glob.glob(search_path)) + if not pt_files: + if is_main: + print(f"No files matching '{search_path}' found. Exiting.") + return + + pt_files_local = pt_files[rank::world_size] + if is_main: + print( + f"Found {len(pt_files)} latent file(s), {world_size} GPU(s), " + f"~{len(pt_files_local)} per GPU.", + flush=True, + ) + + pbar = tqdm(pt_files_local, desc=f"[Rank {rank}] Decoding", disable=(not is_main)) + for pt_path in pbar: + basename = os.path.splitext(os.path.basename(pt_path))[0] + if basename.startswith("latents_"): + video_name = basename.replace("latents_", "video_", 1) + else: + video_name = f"video_{basename}" + out_path = os.path.join(output_dir, f"{video_name}.mp4") + + try: + latent = _load_latent_tensor(pt_path) + except Exception as exc: + print(f"[Rank {rank}] Failed to load {pt_path}: {exc}", flush=True) + continue + + try: + video = decode_latent_to_video(vae, latent, device, dtype) + video_frames = video[0].cpu() + video_uint8 = (video_frames * 255.0).clamp(0, 255).to(torch.uint8) + video_uint8 = video_uint8.permute(0, 2, 3, 1) + write_video(out_path, video_uint8, fps=args.fps) + except Exception as exc: + print(f"[Rank {rank}] Failed to decode {pt_path}: {exc}", flush=True) + + if world_size > 1: + dist.barrier() + if is_main: + print(f"Done. Decoded {len(pt_files)} latent(s) to {output_dir}", flush=True) + if world_size > 1: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/scripts/decode_vae_latents.py b/scripts/decode_vae_latents.py new file mode 100644 index 0000000..03135f2 --- /dev/null +++ b/scripts/decode_vae_latents.py @@ -0,0 +1,167 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0 +# +# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied. +# +# SPDX-License-Identifier: Apache-2.0 + +""" +Decode saved latent files (.pt) to pixel video and write .mp4 using the specified VAE. + +Supports multi-GPU via torchrun: + torchrun --nproc_per_node=8 scripts/decode_vae_latents.py --input_dir /path/to/latents +Single-GPU also works: + python scripts/decode_vae_latents.py --input_dir /path/to/latents + +Uses the Wan2.2-TI2V-5B VAE. +""" +import sys +import os +import argparse +import glob + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import torch +import torch.distributed as dist +from torchvision.io import write_video +from tqdm import tqdm + + +def init_distributed(): + """Initialize distributed process group if launched via torchrun. Returns (rank, world_size).""" + if "RANK" in os.environ and "WORLD_SIZE" in os.environ: + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend="nccl") + return rank, world_size, local_rank + return 0, 1, 0 + + +def get_vae(): + """Return the Wan2.2-TI2V-5B VAE wrapper.""" + from utils.wan_5b_wrapper import WanVAEWrapper + return WanVAEWrapper() + + +def decode_latent_to_video(vae, latent: torch.Tensor, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + """ + Decode latent to pixel video; same logic as pipeline/causal_diffusion_inference.py. + + Args: + vae: Wan2.2-TI2V-5B VAE wrapper + latent: shape (batch, T, C, H, W) or (T, C, H, W) + device, dtype: device and dtype for computation + + Returns: + video: shape (batch, T, C, H, W), range [0, 1] + """ + if latent.dim() == 4: + latent = latent.unsqueeze(0) + latent = latent.to(device=device, dtype=dtype) + video = vae.decode_to_pixel(latent) + video = (video * 0.5 + 0.5).clamp(0, 1) + return video + + +def main(): + parser = argparse.ArgumentParser(description="Decode saved latents to video.") + parser.add_argument( + "--input_dir", + type=str, + required=True, + help="Directory containing .pt latent files (e.g. latents_rank00_idx000000.pt).", + ) + parser.add_argument( + "--output_dir", + type=str, + default=None, + help="Directory to save decoded .mp4 videos. Default: input_dir/decoded_videos", + ) + parser.add_argument( + "--fps", + type=int, + default=24, + help="FPS for output video (default: 16).", + ) + parser.add_argument( + "--pattern", + type=str, + default="latents_*.pt", + help="Glob pattern for latent files (default: latents_*.pt).", + ) + args = parser.parse_args() + + rank, world_size, local_rank = init_distributed() + is_main = (rank == 0) + + output_dir = args.output_dir or os.path.join(args.input_dir, "decoded_videos") + if is_main: + os.makedirs(output_dir, exist_ok=True) + if world_size > 1: + dist.barrier() + + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + torch.set_grad_enabled(False) + device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu") + dtype = torch.bfloat16 + + vae = get_vae() + vae = vae.to(device=device, dtype=dtype).eval() + + search_path = os.path.join(args.input_dir, args.pattern) + pt_files = sorted(glob.glob(search_path)) + if not pt_files: + if is_main: + print(f"No files matching '{search_path}' found. Exiting.") + return + + # Shard files across ranks + pt_files_local = pt_files[rank::world_size] + if is_main: + print(f"Found {len(pt_files)} latent file(s), {world_size} GPU(s), ~{len(pt_files_local)} per GPU.") + + pbar = tqdm(pt_files_local, desc=f"[Rank {rank}] Decoding", disable=(not is_main)) + for pt_path in pbar: + basename = os.path.splitext(os.path.basename(pt_path))[0] + video_name = basename.replace("latents_", "video_", 1) if basename.startswith("latents_") else f"video_{basename}" + out_path = os.path.join(output_dir, f"{video_name}.mp4") + + try: + data = torch.load(pt_path, map_location="cpu", weights_only=True) + if isinstance(data, torch.Tensor): + latent = data + elif isinstance(data, dict): + latent = next(iter(data.values())) + if not isinstance(latent, torch.Tensor): + print(f"[Rank {rank}] Skip {pt_path}: dict value is not a tensor.") + continue + else: + print(f"[Rank {rank}] Skip {pt_path}: unsupported type {type(data)}.") + continue + except Exception as e: + print(f"[Rank {rank}] Failed to load {pt_path}: {e}") + continue + + video = decode_latent_to_video(vae, latent, device, dtype) + video_frames = video[0].cpu() + video_uint8 = (video_frames * 255.0).clamp(0, 255).to(torch.uint8) + video_uint8 = video_uint8.permute(0, 2, 3, 1) # (T, C, H, W) -> (T, H, W, C) + write_video(out_path, video_uint8, fps=args.fps) + + if world_size > 1: + dist.barrier() + if is_main: + print(f"Done. Decoded {len(pt_files)} latent(s) to {output_dir}") + if world_size > 1: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/scripts/merge_lora_generator.py b/scripts/merge_lora_generator.py new file mode 100644 index 0000000..4fe7f2f --- /dev/null +++ b/scripts/merge_lora_generator.py @@ -0,0 +1,130 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0 +# +# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied. +# +# SPDX-License-Identifier: Apache-2.0 +"""Merge a LongLive generator checkpoint with LoRA weights for simple inference.""" + +from __future__ import annotations + +import argparse +import os +import sys +from pathlib import Path + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + + +def _torch_load(path: str): + import torch + + try: + return torch.load(path, map_location="cpu", weights_only=False) + except TypeError: + return torch.load(path, map_location="cpu") + + +def _load_lora_state(path: str): + checkpoint = _torch_load(path) + if isinstance(checkpoint, dict) and "generator_lora" in checkpoint: + return checkpoint["generator_lora"] + return checkpoint + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config_path", required=True, help="Inference yaml containing model, checkpoint, and adapter settings.") + parser.add_argument("--output_path", required=True, help="Path to save the merged generator checkpoint.") + parser.add_argument("--generator_ckpt", default=None, help="Override checkpoints.generator_ckpt from the yaml.") + parser.add_argument("--lora_ckpt", default=None, help="Override checkpoints.lora_ckpt from the yaml.") + parser.add_argument("--device", default="cuda:0", help="Device used for merging, e.g. cuda:0 or cpu.") + parser.add_argument("--dtype", choices=("bf16", "fp32"), default="bf16", help="Save merged weights in this dtype.") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + + import torch + from omegaconf import OmegaConf + + from utils.config import normalize_config + from utils.inference_utils import load_generator_checkpoint + from utils.lora_utils import configure_lora_for_model + from utils.nvfp4_checkpoint import cpu_state_dict + from utils.wan_5b_wrapper import WanDiffusionWrapper + + config = normalize_config(OmegaConf.load(args.config_path)) + generator_ckpt = args.generator_ckpt or getattr(config, "generator_ckpt", None) + lora_ckpt = args.lora_ckpt or getattr(config, "lora_ckpt", None) + if not generator_ckpt: + raise ValueError("Missing generator checkpoint. Set checkpoints.generator_ckpt or pass --generator_ckpt.") + if not lora_ckpt: + raise ValueError("Missing LoRA checkpoint. Set checkpoints.lora_ckpt or pass --lora_ckpt.") + if not getattr(config, "adapter", None): + raise ValueError("Missing adapter config. The merge script needs the LoRA rank/alpha/dropout settings.") + + device = torch.device(args.device) + dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32 + + print(f"Building generator: {config.model_kwargs}") + generator = WanDiffusionWrapper(**getattr(config, "model_kwargs", {}), is_causal=True) + generator.eval().requires_grad_(False) + + print(f"Loading generator checkpoint: {generator_ckpt}") + incompatible = load_generator_checkpoint( + generator, + generator_ckpt, + use_ema=bool(getattr(config, "use_ema", False)), + ) + missing = getattr(incompatible, "missing_keys", []) + unexpected = getattr(incompatible, "unexpected_keys", []) + if missing: + print(f"[Warning] Missing generator keys: {missing[:8]} ...") + if unexpected: + print(f"[Warning] Unexpected generator keys: {unexpected[:8]} ...") + + print(f"Applying LoRA config: {config.adapter}") + generator.model = configure_lora_for_model( + generator.model, + model_name="generator", + lora_config=config.adapter, + is_main_process=True, + ) + + import peft + + print(f"Loading LoRA checkpoint: {lora_ckpt}") + peft.set_peft_model_state_dict(generator.model, _load_lora_state(lora_ckpt)) # type: ignore[arg-type] + + print(f"Merging LoRA on {device} in {dtype}...") + generator.to(device=device, dtype=dtype) + generator.model = generator.model.merge_and_unload(safe_merge=True) + generator.eval().requires_grad_(False) + + output_path = Path(args.output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + checkpoint = { + "generator": cpu_state_dict(generator), + "checkpoint_format": "longlive_generator_merged_lora", + "source_generator_ckpt": str(generator_ckpt), + "source_lora_ckpt": str(lora_ckpt), + "model_name": getattr(config.model_kwargs, "model_name", None), + "dtype": str(dtype).replace("torch.", ""), + "merged_lora": True, + } + torch.save(checkpoint, output_path) + size_gib = os.path.getsize(output_path) / (1024 ** 3) + print(f"Saved merged generator to {output_path} ({size_gib:.2f} GiB).") + print("Use this file as checkpoints.generator_ckpt for inference and remove adapter/lora_ckpt from the inference config.") + + +if __name__ == "__main__": + main() diff --git a/scripts/save_merged_nvfp4_generator.py b/scripts/save_merged_nvfp4_generator.py new file mode 100644 index 0000000..8be1007 --- /dev/null +++ b/scripts/save_merged_nvfp4_generator.py @@ -0,0 +1,246 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0 +# +# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied. +# +# SPDX-License-Identifier: Apache-2.0 +"""Save a merged LoRA generator as a reusable checkpoint.""" +from __future__ import annotations + +import argparse +import gc +import os +import sys +from pathlib import Path + +import torch +from omegaconf import OmegaConf + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from utils.config import normalize_config +from utils.lora_utils import configure_lora_for_model +from utils.nvfp4_checkpoint import ( + NVFP4_CHECKPOINT_FORMAT, + NVFP4_CHECKPOINT_VERSION, + clean_fsdp_state_dict_keys, + cpu_state_dict, + is_nvfp4_state_dict, + quantize_model_for_fouroversix_nvfp4, + unwrap_generator_state_dict, +) +from utils.quant import _materialize_quantized_weights_for_inference +from utils.wan_5b_wrapper import WanDiffusionWrapper + + +def _torch_load(path: str): + try: + return torch.load(path, map_location="cpu", weights_only=False) + except TypeError: + return torch.load(path, map_location="cpu") + + +def _load_generator_checkpoint(generator: WanDiffusionWrapper, checkpoint_path: str, use_ema: bool) -> None: + checkpoint = _torch_load(checkpoint_path) + state_dict = unwrap_generator_state_dict(checkpoint, use_ema=use_ema) + if is_nvfp4_state_dict(state_dict): + raise ValueError( + f"{checkpoint_path} is already a materialized NVFP4 checkpoint; " + "use it directly as checkpoints.generator_ckpt." + ) + if use_ema: + state_dict = clean_fsdp_state_dict_keys(state_dict) + incompatible = generator.load_state_dict(state_dict, strict=not use_ema) + missing = getattr(incompatible, "missing_keys", []) + unexpected = getattr(incompatible, "unexpected_keys", []) + if missing: + print(f"[Warning] Missing generator keys while loading base checkpoint: {missing[:8]} ...") + if unexpected: + print(f"[Warning] Unexpected generator keys while loading base checkpoint: {unexpected[:8]} ...") + + +def _load_lora_state(lora_ckpt_path: str): + checkpoint = _torch_load(lora_ckpt_path) + if isinstance(checkpoint, dict) and "generator_lora" in checkpoint: + return checkpoint["generator_lora"] + return checkpoint + + +def _merge_lora(generator: WanDiffusionWrapper, config, lora_ckpt_path: str) -> WanDiffusionWrapper: + if not getattr(config, "adapter", None): + raise ValueError("LoRA merge was requested, but config.adapter is missing.") + if not lora_ckpt_path: + raise ValueError("LoRA merge was requested, but no lora_ckpt was provided.") + + print(f"Applying LoRA config: {config.adapter}") + generator.model = configure_lora_for_model( + generator.model, + model_name="generator", + lora_config=config.adapter, + is_main_process=True, + ) + + import peft + + print(f"Loading LoRA weights: {lora_ckpt_path}") + peft.set_peft_model_state_dict(generator.model, _load_lora_state(lora_ckpt_path)) # type: ignore[arg-type] + print("Merging LoRA into generator...") + generator.model = generator.model.merge_and_unload(safe_merge=True) + return generator + + +def _metadata( + config, + args: argparse.Namespace, + *, + backend: str, + matched_modules: list[str], + materialized_modules: list[str], +): + checkpoint_format = ( + "longlive_generator_merged_bf16" + if backend == "transformer_engine" + else NVFP4_CHECKPOINT_FORMAT + ) + quant_format = "bf16" if backend == "transformer_engine" else "nvfp4" + quant_backend = "transformer_engine_runtime" if backend == "transformer_engine" else "fouroversix" + return { + "checkpoint_format": checkpoint_format, + "checkpoint_version": NVFP4_CHECKPOINT_VERSION, + "source_generator_ckpt": args.generator_ckpt, + "source_lora_ckpt": args.lora_ckpt, + "merged_lora": bool(args.lora_ckpt and not args.no_merge_lora), + "model_name": getattr(config.model_kwargs, "model_name", None), + "quantization": { + "format": quant_format, + "backend": quant_backend, + "materialized": backend == "fouroversix", + "dtype": "bfloat16", + "scale_rule": getattr(config, "model_quant_scale_rule", "static_6"), + "activation_scale_rule": getattr(config, "model_quant_activation_scale_rule", None), + "weight_scale_rule": getattr(config, "model_quant_weight_scale_rule", None), + "gradient_scale_rule": getattr(config, "model_quant_gradient_scale_rule", None), + "te_inference_only": bool( + getattr(config, "model_quant_te_inference_only", backend == "transformer_engine") + ), + "te_low_precision_weights": bool( + getattr(config, "model_quant_te_low_precision_weights", backend == "transformer_engine") + ), + "te_fallback_to_fouroversix": bool( + getattr(config, "model_quant_te_fallback_to_fouroversix", False) + ), + "matched_filtered_modules": matched_modules, + "materialized_modules": materialized_modules, + }, + } + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Merge generator LoRA weights and save either packed FourOverSix NVFP4 or TE-ready bf16." + ) + parser.add_argument("--config_path", required=True, help="Inference yaml that contains model/adapter/quant settings.") + parser.add_argument("--output_path", required=True, help="Path to save the generator .pt file.") + parser.add_argument("--generator_ckpt", default=None, help="Override checkpoints.generator_ckpt from the yaml.") + parser.add_argument("--lora_ckpt", default=None, help="Override checkpoints.lora_ckpt from the yaml.") + parser.add_argument("--device", default="cuda:0", help="Device used for quantization, e.g. cuda:0 or cpu.") + parser.add_argument( + "--backend", + choices=("fouroversix", "transformer_engine"), + default="fouroversix", + help=( + "fouroversix saves packed/materialized NVFP4. " + "transformer_engine saves merged bf16 for TE runtime quantization." + ), + ) + parser.add_argument("--no_merge_lora", action="store_true", help="Quantize the base generator without merging LoRA.") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + config = normalize_config(OmegaConf.load(args.config_path)) + args.generator_ckpt = args.generator_ckpt or getattr(config, "generator_ckpt", None) + args.lora_ckpt = args.lora_ckpt or getattr(config, "lora_ckpt", None) + + if not args.generator_ckpt: + raise ValueError("Missing generator checkpoint. Set checkpoints.generator_ckpt or pass --generator_ckpt.") + config.model_quant_use_transformer_engine = args.backend == "transformer_engine" + + device = torch.device(args.device) + print(f"Building generator on CPU: {config.model_kwargs}") + generator = WanDiffusionWrapper(**getattr(config, "model_kwargs", {}), is_causal=True) + generator.eval().requires_grad_(False) + + print(f"Loading base generator checkpoint: {args.generator_ckpt}") + _load_generator_checkpoint(generator, args.generator_ckpt, use_ema=bool(getattr(config, "use_ema", False))) + + should_merge_lora = bool(getattr(config, "merge_lora", False)) and not args.no_merge_lora + if should_merge_lora: + generator = _merge_lora(generator, config, args.lora_ckpt) + else: + print("Skipping LoRA merge; quantizing the loaded generator as-is.") + + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + print(f"Moving generator to {device} and casting to bfloat16...") + generator.to(device=device, dtype=torch.bfloat16) + materialized_modules = [] + if args.backend == "transformer_engine": + matched_modules = [] + print( + "Saving merged bf16 weights for TransformerEngine runtime quantization. " + "TransformerEngine state_dict is not a packed NVFP4 storage format." + ) + else: + generator.model, matched_modules = quantize_model_for_fouroversix_nvfp4( + generator.model, + config=config, + keep_master_weights=False, + verbose=True, + ) + + print("Materializing NVFP4 weights and dropping bf16 master weights...") + materialized_modules, master_bytes, quantized_bytes = _materialize_quantized_weights_for_inference( + generator.model, + target_device=device, + ) + print( + "[NVFP4] " + f"materialized_modules={len(materialized_modules)}, " + f"master_weight={master_bytes / (1024 ** 3):.3f} GiB, " + f"quantized_weight={quantized_bytes / (1024 ** 3):.3f} GiB" + ) + + print("Copying checkpoint tensors to CPU for saving...") + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + output_path = Path(args.output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + checkpoint = { + "generator": cpu_state_dict(generator), + **_metadata( + config, + args, + backend=args.backend, + matched_modules=matched_modules, + materialized_modules=materialized_modules, + ), + } + torch.save(checkpoint, output_path) + size_gib = os.path.getsize(output_path) / (1024 ** 3) + print(f"Saved {args.backend} generator checkpoint to {output_path} ({size_gib:.3f} GiB)") + + +if __name__ == "__main__": + with torch.no_grad(): + main() diff --git a/tests/test_dmd_i2v_conditioning.py b/tests/test_dmd_i2v_conditioning.py new file mode 100644 index 0000000..ae142a4 --- /dev/null +++ b/tests/test_dmd_i2v_conditioning.py @@ -0,0 +1,51 @@ +import unittest + +import torch + +from utils.i2v_conditioning import ( + _get_i2v_context_frames, + _i2v_loss_mask_like, + _overwrite_i2v_context, +) + + +class DMDI2VConditioningTest(unittest.TestCase): + def test_overwrite_i2v_context_keeps_only_initial_frames_clean(self): + tensor = torch.randn(2, 4, 3, 2, 2) + original = tensor.clone() + initial_latent = torch.full((2, 1, 3, 2, 2), 7.0) + + context_frames = _get_i2v_context_frames(tensor, initial_latent) + conditioned = _overwrite_i2v_context(tensor, initial_latent, context_frames) + + self.assertEqual(context_frames, 1) + self.assertTrue(torch.equal(conditioned[:, :1], initial_latent)) + self.assertTrue(torch.equal(conditioned[:, 1:], original[:, 1:])) + + def test_i2v_loss_mask_excludes_context_frames(self): + tensor = torch.randn(2, 4, 3, 2, 2) + mask = _i2v_loss_mask_like(tensor, context_frames=1) + + self.assertFalse(mask[:, :1].any()) + self.assertTrue(mask[:, 1:].all()) + + def test_overwrite_i2v_context_keeps_chunk_length_unchanged(self): + initial_latent = torch.full((2, 1, 3, 2, 2), 5.0) + noise_chunk = torch.randn(2, 4, 3, 2, 2) + + first_chunk = _overwrite_i2v_context(noise_chunk, initial_latent, context_frames=1) + + self.assertEqual(first_chunk.shape[1], 4) + self.assertTrue(torch.equal(first_chunk[:, :1], initial_latent)) + self.assertTrue(torch.equal(first_chunk[:, 1:], noise_chunk[:, 1:])) + + def test_i2v_context_must_not_cover_entire_clip(self): + tensor = torch.randn(2, 1, 3, 2, 2) + initial_latent = torch.randn(2, 1, 3, 2, 2) + + with self.assertRaises(ValueError): + _get_i2v_context_frames(tensor, initial_latent) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_fp8_inference.py b/tests/test_fp8_inference.py new file mode 100644 index 0000000..2c3cdb1 --- /dev/null +++ b/tests/test_fp8_inference.py @@ -0,0 +1,29 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +import torch.nn as nn + +from utils.fp8 import quantize_model_fp8 +from utils.torch_compile_utils import SafeCompiledCallable + + +def test_fp8_quantization_requires_cuda(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + + with pytest.raises(RuntimeError, match="requires a CUDA GPU"): + quantize_model_fp8(nn.Linear(16, 16, dtype=torch.bfloat16)) + + +def test_strict_compile_wrapper_does_not_shadow_torch(monkeypatch): + monkeypatch.setattr(torch, "compile", lambda fn, **_kwargs: fn) + + compiled = SafeCompiledCallable( + lambda value: value + 1, + name="test", + mode=None, + suppress_errors=False, + ) + + assert compiled(1) == 2 diff --git a/tests/test_i2v_dataset_frame_accounting.py b/tests/test_i2v_dataset_frame_accounting.py new file mode 100644 index 0000000..c8c80d6 --- /dev/null +++ b/tests/test_i2v_dataset_frame_accounting.py @@ -0,0 +1,40 @@ +import tempfile +import unittest +from pathlib import Path + +from utils.dataset import MultiVideoConcatDataset + + +class I2VDatasetFrameAccountingTest(unittest.TestCase): + def _dataset(self, root, latent_frames): + (root / "video" / "sample").mkdir(parents=True) + (root / "caption" / "sample").mkdir(parents=True) + total_raw_frames = 1 + (latent_frames - 1) * 4 + return MultiVideoConcatDataset( + data_dir=str(root), + video_size=(704, 1280), + total_frames=total_raw_frames, + independent_first_frame=True, + num_frame_per_block=8, + temporal_compression_ratio=4, + ) + + def test_96_frame_i2v_uses_regular_8_frame_blocks(self): + with tempfile.TemporaryDirectory() as tmp: + dataset = self._dataset(Path(tmp), latent_frames=96) + + self.assertEqual(dataset.first_chunk_latent_frames, 8) + self.assertEqual(dataset.first_chunk_frames, 29) + self.assertEqual(dataset.total_segments, 12) + + def test_33_frame_i2v_keeps_legacy_one_plus_block_layout(self): + with tempfile.TemporaryDirectory() as tmp: + dataset = self._dataset(Path(tmp), latent_frames=33) + + self.assertEqual(dataset.first_chunk_latent_frames, 9) + self.assertEqual(dataset.first_chunk_frames, 33) + self.assertEqual(dataset.total_segments, 4) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_i2v_sequence_parallel_config.py b/tests/test_i2v_sequence_parallel_config.py new file mode 100644 index 0000000..c356a90 --- /dev/null +++ b/tests/test_i2v_sequence_parallel_config.py @@ -0,0 +1,63 @@ +import unittest +from types import SimpleNamespace + +from wan_5b.distributed.sp_training import ( + sp_training_sequence_frame_count, + validate_sequence_parallel_training_config, +) + + +class I2VSequenceParallelConfigTest(unittest.TestCase): + def test_i2v_training_frames_use_full_configured_sequence(self): + cfg = SimpleNamespace( + i2v=True, + independent_first_frame=True, + image_or_video_shape=[1, 96, 48, 44, 80], + ) + + self.assertEqual(sp_training_sequence_frame_count(cfg), 96) + + def test_i2v_sequence_parallel_partitions_the_96_frame_sequence(self): + cfg = SimpleNamespace( + i2v=True, + independent_first_frame=True, + image_or_video_shape=[1, 96, 48, 44, 80], + ) + + validate_sequence_parallel_training_config( + cfg, + sp_size=4, + num_frame_per_block=8, + ) + + def test_i2v_sequence_parallel_rejects_non_block_aligned_length(self): + cfg = SimpleNamespace( + i2v=True, + independent_first_frame=True, + image_or_video_shape=[1, 97, 48, 44, 80], + ) + + with self.assertRaisesRegex(ValueError, r"training latent frames"): + validate_sequence_parallel_training_config( + cfg, + sp_size=4, + num_frame_per_block=8, + ) + + def test_t2v_sequence_parallel_uses_full_sequence(self): + cfg = SimpleNamespace( + i2v=False, + independent_first_frame=False, + image_or_video_shape=[1, 96, 48, 44, 80], + ) + + self.assertEqual(sp_training_sequence_frame_count(cfg), 96) + validate_sequence_parallel_training_config( + cfg, + sp_size=4, + num_frame_per_block=8, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_i2v_teacher_forcing_context.py b/tests/test_i2v_teacher_forcing_context.py new file mode 100644 index 0000000..c1ba093 --- /dev/null +++ b/tests/test_i2v_teacher_forcing_context.py @@ -0,0 +1,192 @@ +import unittest +import importlib.util +import pathlib +import sys +import types +from types import SimpleNamespace + +import torch + + +class _FakeScheduler: + def __init__(self): + self.num_train_timesteps = 1000 + self.timesteps = torch.arange(1000, dtype=torch.float32) + self.sigmas = torch.linspace(1.0, 0.0, 1000) + + def add_noise(self, clean, noise, timestep): + return clean + 10.0 + + def training_target(self, clean, noise, timestep): + return torch.zeros_like(clean) + + def training_weight(self, timestep): + return torch.ones(timestep.numel(), device=timestep.device, dtype=torch.float32) + + +class _FakeGenerator: + def __init__(self): + self.recorded = None + + def __call__( + self, + *, + noisy_image_or_video, + conditional_dict, + timestep, + clean_x, + aug_t, + ): + self.recorded = { + "noisy_image_or_video": noisy_image_or_video.detach().clone(), + "timestep": timestep.detach().clone(), + "clean_x": clean_x.detach().clone(), + "aug_t": aug_t.detach().clone() if aug_t is not None else None, + } + return torch.zeros_like(noisy_image_or_video), torch.zeros_like(noisy_image_or_video) + + +class _FakeBuffer: + num_blocks = 0 + + def is_empty(self): + return False + + def add(self, error_block, timestep_index, block_pos=None): + pass + + def stats(self): + return { + "total_added": 0, + "filled_buckets": "0/0", + "total_entries": 0, + } + + +class _StubBaseModel(torch.nn.Module): + def _get_timestep( + self, + min_timestep, + max_timestep, + batch_size, + num_frame, + num_frame_per_block, + uniform_timestep=False, + ): + if uniform_timestep: + return torch.randint( + min_timestep, + max_timestep, + [batch_size, 1], + device=self.device, + dtype=torch.long, + ).repeat(1, num_frame) + timestep = torch.randint( + min_timestep, + max_timestep, + [batch_size, num_frame], + device=self.device, + dtype=torch.long, + ) + timestep = timestep.reshape(timestep.shape[0], -1, num_frame_per_block) + timestep[:, :, 1:] = timestep[:, :, 0:1] + return timestep.reshape(timestep.shape[0], -1) + + +def _load_causal_diffusion_with_stubs(): + repo_root = pathlib.Path(__file__).resolve().parents[1] + module_path = repo_root / "model" / "diffusion.py" + saved = { + name: sys.modules.get(name) + for name in ( + "model", + "model.base", + "pipeline", + "utils.wan_5b_wrapper", + ) + } + + fake_model = types.ModuleType("model") + fake_model.__path__ = [] + fake_base = types.ModuleType("model.base") + fake_base.BaseModel = _StubBaseModel + + fake_pipeline = types.ModuleType("pipeline") + fake_pipeline.CausalDiffusionInferencePipeline = object + + fake_wrapper = types.ModuleType("utils.wan_5b_wrapper") + fake_wrapper.WanDiffusionWrapper = object + fake_wrapper.WanTextEncoder = object + fake_wrapper.WanVAEWrapper = object + + sys.modules["model"] = fake_model + sys.modules["model.base"] = fake_base + sys.modules["pipeline"] = fake_pipeline + sys.modules["utils.wan_5b_wrapper"] = fake_wrapper + try: + spec = importlib.util.spec_from_file_location( + "_diffusion_under_test", module_path + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.CausalDiffusion + finally: + for name, value in saved.items(): + if value is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = value + + +class I2VTeacherForcingContextTest(unittest.TestCase): + def test_teacher_forcing_keeps_i2v_context_clean_after_augmentation(self): + CausalDiffusion = _load_causal_diffusion_with_stubs() + model = CausalDiffusion.__new__(CausalDiffusion) + torch.nn.Module.__init__(model) + model.device = torch.device("cpu") + model.dtype = torch.float32 + model.args = SimpleNamespace(i2v=True) + model.independent_first_frame = True + model.num_frame_per_block = 2 + model.scheduler = _FakeScheduler() + model.teacher_forcing = True + model.noise_augmentation_max_timestep = 10 + model.generator = _FakeGenerator() + model.error_buffer = _FakeBuffer() + model.noise_error_buffer = _FakeBuffer() + model.er_start_step = 0 + model.er_clean_prob = 0.0 + model.er_latent_inject_prob = 0.0 + model.er_noise_inject_prob = 0.0 + model.er_context_inject_prob = 1.0 + model.er_buffer_warmup_iter = -1 + + def inject_context_error(clean_latent_aug, index, batch_size, num_frame): + return clean_latent_aug + 100.0 + + model._inject_error_buffer = inject_context_error + + clean_latent = torch.zeros(1, 4, 1, 1, 1) + initial_latent = torch.full((1, 1, 1, 1, 1), 7.0) + + model.generator_loss( + image_or_video_shape=[1, 4, 1, 1, 1], + conditional_dict={"prompt_embeds": torch.zeros(1, 1)}, + unconditional_dict={}, + clean_latent=clean_latent, + initial_latent=initial_latent, + global_step=0, + ) + + recorded = model.generator.recorded + self.assertTrue(torch.equal(recorded["noisy_image_or_video"][:, :1], initial_latent)) + self.assertTrue(torch.equal(recorded["clean_x"][:, :1], initial_latent)) + self.assertTrue((recorded["timestep"][:, :1] == 0).all()) + self.assertTrue((recorded["aug_t"][:, :1] == 0).all()) + + # Later frames still receive clean-side augmentation and context error. + self.assertTrue(torch.equal(recorded["clean_x"][:, 1:], torch.full((1, 3, 1, 1, 1), 110.0))) + + +if __name__ == "__main__": + unittest.main() diff --git a/train.py b/train.py new file mode 100644 index 0000000..2e4f32d --- /dev/null +++ b/train.py @@ -0,0 +1,47 @@ +# Adopted from https://github.com/guandeh17/Self-Forcing +# SPDX-License-Identifier: Apache-2.0 +import argparse +import os +from omegaconf import OmegaConf +import wandb + +from trainer import ScoreDistillationTrainer, DiffusionTrainer +from utils.config import normalize_config + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--config_path", type=str, required=True) + parser.add_argument("--no_save", action="store_true") + parser.add_argument("--no_visualize", action="store_true") + parser.add_argument("--logdir", type=str, default="", help="Path to the directory to save logs") + parser.add_argument("--wandb-save-dir", type=str, default="", help="Path to the directory to save wandb logs") + parser.add_argument("--disable-wandb", action="store_true") + parser.add_argument("--no-auto-resume", action="store_true", help="Disable auto resume from latest checkpoint in logdir") + parser.add_argument("--generate-before-train", action="store_true", help="Run one evaluation inference before training starts") + + args = parser.parse_args() + + config = normalize_config(OmegaConf.load(args.config_path)) + config.no_save = args.no_save + config.no_visualize = args.no_visualize + + config_name = os.path.splitext(os.path.basename(args.config_path))[0] + config.config_name = config_name + config.logdir = args.logdir + config.wandb_save_dir = args.wandb_save_dir + config.disable_wandb = args.disable_wandb + config.auto_resume = not args.no_auto_resume # Default to True unless --no-auto-resume is specified + config.generate_before_train = args.generate_before_train + + if config.trainer == "score_distillation": + trainer = ScoreDistillationTrainer(config) + elif config.trainer == "diffusion": + trainer = DiffusionTrainer(config) + trainer.train() + + wandb.finish() + + +if __name__ == "__main__": + main() diff --git a/trainer/__init__.py b/trainer/__init__.py new file mode 100644 index 0000000..a115e45 --- /dev/null +++ b/trainer/__init__.py @@ -0,0 +1,7 @@ +from .distillation import Trainer as ScoreDistillationTrainer +from .diffusion import Trainer as DiffusionTrainer + +__all__ = [ + "ScoreDistillationTrainer", + "DiffusionTrainer", +] diff --git a/trainer/diffusion.py b/trainer/diffusion.py new file mode 100644 index 0000000..dddcae1 --- /dev/null +++ b/trainer/diffusion.py @@ -0,0 +1,1120 @@ +# Adopted from https://github.com/guandeh17/Self-Forcing +# SPDX-License-Identifier: Apache-2.0 + +import gc +import logging +import types + +from model import CausalDiffusion +from wan_5b.distributed.sp_training import SequenceParallelHelper +from utils.dataset import MultiVideoConcatDataset, MultiTextConcatDataset, cycle, multi_video_collate_fn, eval_collate_fn +from utils.config import section_get, wan_default_config +from utils.misc import set_seed +import torch.distributed as dist +from omegaconf import OmegaConf +import torch +import wandb +import time +import os +from torchvision.io import write_video +from utils.distributed import EMA_FSDP, barrier, fsdp_wrap, launch_distributed_job, FSDP +from torch.distributed.fsdp import ( + StateDictType, FullStateDictConfig, FullOptimStateDictConfig +) + +def save_prompts_to_txt(prompts_for_sample, prompt_txt_path: str, is_main_process: bool): + """ + Save prompts for one generated video to a txt file. + Consecutive identical prompts are merged, e.g.: + [0] a, [1] a, [2] b => [0,1] a\n[2] b\n + """ + try: + with open(prompt_txt_path, "w", encoding="utf-8") as f: + if len(prompts_for_sample) == 0: + return + + current_prompt = prompts_for_sample[0] + current_indices = [0] + for seg_idx in range(1, len(prompts_for_sample)): + p = prompts_for_sample[seg_idx] + if p == current_prompt: + current_indices.append(seg_idx) + else: + indices_str = ",".join(str(i) for i in current_indices) + f.write(f"[{indices_str}] {current_prompt}\n") + current_prompt = p + current_indices = [seg_idx] + # flush the last run + indices_str = ",".join(str(i) for i in current_indices) + f.write(f"[{indices_str}] {current_prompt}\n") + except Exception as e: + if is_main_process: + print(f"Warning: failed to save prompts to {prompt_txt_path}: {e}") + + +class Trainer: + def __init__(self, config): + self.config = config + self.step = 0 + + # Step 1: Initialize the distributed training environment (rank, seed, dtype, logging etc.) + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + + launch_distributed_job() + global_rank = dist.get_rank() + + self.dtype = torch.bfloat16 if config.mixed_precision else torch.float32 + self.device = torch.cuda.current_device() + self.is_main_process = global_rank == 0 + self.causal = config.causal + self.disable_wandb = config.disable_wandb + + # use a random seed for the training + if config.seed == 0: + random_seed = torch.randint(0, 10000000, (1,), device=self.device) + dist.broadcast(random_seed, src=0) + config.seed = random_seed.item() + + set_seed(config.seed + global_rank) + + if self.is_main_process and not self.disable_wandb: + if getattr(config, "wandb_key", None): + wandb.login(host=config.wandb_host, key=config.wandb_key) + wandb.init( + config=OmegaConf.to_container(config, resolve=True), + name=config.config_name, + mode="online", + entity=config.wandb_entity, + project=config.wandb_project, + dir=config.wandb_save_dir + ) + + self.output_path = config.logdir + auto_resume = getattr(config, "auto_resume", True) + self.gradient_accumulation_steps = getattr(config, "gradient_accumulation_steps", 1) + + # Sequence Parallel is supported only for the 5B model; world_size must + # equal sp_size * dp_size. + self.sequence_parallel_size = getattr(config, "sequence_parallel_size", 1) + world_size = dist.get_world_size() + self.data_parallel_size = world_size // self.sequence_parallel_size if self.sequence_parallel_size > 1 else world_size + self.sp_group = None + self.dp_group = None + + if self.is_main_process and self.gradient_accumulation_steps > 1: + eff_batch = config.batch_size * self.gradient_accumulation_steps * self.data_parallel_size + print(f"Gradient accumulation steps: {self.gradient_accumulation_steps}, effective batch size: {eff_batch}") + + if self.sequence_parallel_size > 1: + assert config.model_kwargs.model_name == "Wan2.2-TI2V-5B", ( + f"sequence_parallel_size is only supported for Wan2.2-TI2V-5B model, but got {config.model_kwargs.model_name}" + ) + assert world_size % self.sequence_parallel_size == 0, ( + f"world_size ({world_size}) must be divisible by sequence_parallel_size ({self.sequence_parallel_size})" + ) + from wan_5b.distributed.sp_training import ( + validate_sequence_parallel_training_config, + ) + validate_sequence_parallel_training_config( + config, + self.sequence_parallel_size, + config.num_frame_per_block, + ) + # Create SP process groups: each DP group contains sp_size ranks, + # and all_to_all runs only within that group. + from wan_5b.distributed.sp_training import ( + set_data_parallel_group, + set_sequence_parallel_group, + ) + sp_size = self.sequence_parallel_size + dp_size = self.data_parallel_size + sp_groups = [] + for g in range(dp_size): + ranks_g = list(range(g * sp_size, (g + 1) * sp_size)) + sp_groups.append(dist.new_group(ranks=ranks_g)) + self.sp_group = sp_groups[global_rank // sp_size] + set_sequence_parallel_group(self.sp_group) + + # Also create DP groups: ranks with the same SP rank across DP + # replicas own the same sequence chunk. For sp_rank=k, the DP group + # is [k, sp+k, 2*sp+k, ..., (dp-1)*sp+k]. This lets warmup gather + # different batches of errors for the same block efficiently. + dp_groups = [] + for k in range(sp_size): + ranks_k = [g * sp_size + k for g in range(dp_size)] + dp_groups.append(dist.new_group(ranks=ranks_k)) + self.dp_group = dp_groups[global_rank % sp_size] + set_data_parallel_group(self.dp_group) + if self.is_main_process: + print(f"[SP] Sequence Parallel enabled, sp_size={sp_size}, dp_size={dp_size}, world_size={world_size}") + + # Step 2: Initialize the model and optimizer + self.model = CausalDiffusion(config, device=self.device) + self.sp_helper = SequenceParallelHelper(self) + + # 2D mode only: print which GLOBAL block-position slice this rank is + # responsible for. The LAST SP rank carries the most error-accumulated + # tail blocks, useful when debugging position-bucketed error recycling. + if self.model.error_buffer is not None and self.model.er_num_blocks > 0: + lo = self.model.er_block_offset + hi = lo + self.model.er_num_blocks + global_rank_id = dist.get_rank() + sp_rk = global_rank_id % max(self.sequence_parallel_size, 1) + print( + f"[ErrorBuffer] rank={global_rank_id} sp_rank={sp_rk} " + f"covers GLOBAL blocks [{lo},{hi}) ({self.model.er_num_blocks} local blocks)" + ) + + # Bind the SP forward path before FSDP wrapping. + model_name = getattr(getattr(config, "model_kwargs", None), "model_name", "") or "" + if self.sequence_parallel_size > 1 and "Wan2.2-TI2V-5B" in model_name: + from wan_5b.distributed.sequence_parallel import ( + sp_dit_causal_forward_train, + sp_causal_attn_forward, + ) + model = self.model.generator.model + # Use the SP forward implementation in the training path. + model._forward_train = types.MethodType(sp_dit_causal_forward_train, model) + + # Keep the original self_attn.forward so inference can temporarily + # disable SP. + self._sp_attn_blocks = [] + for block in model.blocks: + sa = block.self_attn + if not hasattr(sa, "_orig_forward"): + sa._orig_forward = sa.forward + sa.forward = types.MethodType(sp_causal_attn_forward, sa) + self._sp_attn_blocks.append(sa) + + if self.is_main_process: + print("[SP] sp_dit_causal_forward_train and sp_causal_attn_forward are enabled") + print("[SP] natural TF layout is the default training layout") + if getattr(config, "load_raw_video", False): + print(f"[SP-VAE] chunk-halo VAE enabled, halo_latents={self.sp_helper.vae_halo_latents}") + + # ================================= NVFP4 Quantized Training ================================= + self.model_quant = getattr(config, "model_quant", False) + if self.model_quant: + from utils.quant import ModelQuantizationConfig, quantize_model_with_filter + + quant_cfg = ModelQuantizationConfig( + scale_rule=getattr(config, "model_quant_scale_rule", "static_6"), + activation_scale_rule=getattr(config, "model_quant_activation_scale_rule", "static_6"), + weight_scale_rule=getattr(config, "model_quant_weight_scale_rule", None), + gradient_scale_rule=getattr(config, "model_quant_gradient_scale_rule", None), + keep_master_weights=True, + weight_scale_2d=True, + ) + self.model.generator.model, matched_modules = quantize_model_with_filter( + self.model.generator.model, + quant_config=quant_cfg, + filtered_modules=getattr(config, "model_quant_filtered_modules", None), + use_default_filtered_modules=getattr(config, "model_quant_use_default_filtered_modules", True), + cast_model_to_bf16=False, + materialize_for_inference=False, + verbose=self.is_main_process, + ) + if self.is_main_process: + from fouroversix.matmul.cutlass.backend import CUTLASSMatmulBackend + + print(f"[NVFP4] CUTLASS available: {CUTLASSMatmulBackend.is_available()}") + print( + "[NVFP4] Quantized AR training enabled " + "(keep_master_weights=True, weight_scale_2d=True)" + ) + print(f"[NVFP4] {len(matched_modules)} modules excluded from quantization") + + # ================================= Load model weights (before FSDP) ================================= + # Load model weights before FSDP wrapping, while keys still match the + # raw nn.Module. Optimizer, EMA, and step state are restored after FSDP + # and the related objects are created, so keep raw_state. + # + # Priority: auto_resume from logdir > generator_ckpt for a + # cold start > random initialization. This allows configs to keep + # generator_ckpt set while interrupted training still resumes from the + # latest step. The style mirrors trainer/distillation.py. + raw_state = None + + checkpoint_path = None + + if auto_resume and self.output_path: + latest_checkpoint = self.find_latest_checkpoint(self.output_path) + if latest_checkpoint: + checkpoint_path = latest_checkpoint + if self.is_main_process: + print(f"Auto resume: Found latest checkpoint at {checkpoint_path}") + else: + if self.is_main_process: + print("Auto resume: No checkpoint found in logdir, starting from scratch") + elif auto_resume: + if self.is_main_process: + print("Auto resume enabled but no logdir specified, starting from scratch") + else: + if self.is_main_process: + print("Auto resume disabled, starting from scratch") + + if checkpoint_path is None and getattr(config, "generator_ckpt", False): + checkpoint_path = config.generator_ckpt + if self.is_main_process: + print(f"Using explicit checkpoint: {checkpoint_path}") + + if checkpoint_path: + if self.is_main_process: + print(f"Loading checkpoint from {checkpoint_path}") + checkpoint = torch.load(checkpoint_path, map_location="cpu") + + if "generator" in checkpoint: + if self.is_main_process: + print(f"Loading pretrained generator from {checkpoint_path}") + self.model.generator.load_state_dict(checkpoint["generator"], strict=True) + del checkpoint["generator"] + elif "model" in checkpoint: + if self.is_main_process: + print(f"Loading pretrained generator from {checkpoint_path}") + self.model.generator.load_state_dict(checkpoint["model"], strict=True) + del checkpoint["model"] + else: + if self.is_main_process: + print(f"No 'generator'/'model' key found in {checkpoint_path}, treating as raw state_dict") + self.model.generator.load_state_dict(checkpoint, strict=True) + + gc.collect() + + raw_state = checkpoint + if "step" in raw_state: + self.step = raw_state["step"] + if self.is_main_process: + print(f"Resuming from step {self.step}") + else: + if self.is_main_process: + print("Warning: Step not found in checkpoint, starting from step 0.") + + # ================================= FSDP Wrap ================================= + self.model.generator = fsdp_wrap( + self.model.generator, + sharding_strategy=config.sharding_strategy, + mixed_precision=config.mixed_precision, + wrap_strategy=config.generator_fsdp_wrap_strategy + ) + + self.model.text_encoder = fsdp_wrap( + self.model.text_encoder, + sharding_strategy=config.sharding_strategy, + mixed_precision=config.mixed_precision, + wrap_strategy=config.text_encoder_fsdp_wrap_strategy + ) + + if not config.no_visualize or config.load_raw_video: + self.model.vae = self.model.vae.to( + device=self.device, dtype=torch.bfloat16 if config.mixed_precision else torch.float32) + + rename_param = ( + lambda name: name.replace("_fsdp_wrapped_module.", "") + .replace("_checkpoint_wrapped_module.", "") + .replace("_orig_mod.", "") + ) + self.name_to_trainable_params = {} + for n, p in self.model.generator.named_parameters(): + if not p.requires_grad: + continue + + renamed_n = rename_param(n) + self.name_to_trainable_params[renamed_n] = p + + self.generator_optimizer = torch.optim.AdamW( + [param for param in self.model.generator.parameters() + if param.requires_grad], + lr=config.lr, + betas=(config.beta1, config.beta2), + weight_decay=config.weight_decay + ) + + # Step 3: Initialize the dataloader + frame_raw_height = list(config.image_or_video_shape)[3] * wan_default_config[config.model_kwargs.model_name]["spatial_compression_ratio"] + frame_raw_width = list(config.image_or_video_shape)[4] * wan_default_config[config.model_kwargs.model_name]["spatial_compression_ratio"] + total_frames = (list(config.image_or_video_shape)[1] - 1) * wan_default_config[config.model_kwargs.model_name]["temporal_compression_ratio"] + 1 + num_frame_per_block = config.num_frame_per_block + self.fps = wan_default_config[config.model_kwargs.model_name].get("fps", 16) + + allow_padding = getattr(config, "allow_padding", False) + min_latent_frames = getattr(config, "min_latent_frames", 0) + single_video_only = getattr(config, "uniform_prompt", False) + max_chunks_per_shot = getattr(config, "max_chunks_per_shot", 0) + dataset_sample_warning_seconds = getattr(config, "dataset_sample_warning_seconds", 60.0) + dataset_sample_warning_interval_seconds = getattr( + config, "dataset_sample_warning_interval_seconds", 60.0 + ) + dataset = MultiVideoConcatDataset( + data_dir=config.data_path, + video_size=(frame_raw_height, frame_raw_width), + total_frames=total_frames, + deterministic=False, + num_frame_per_block=num_frame_per_block, + temporal_compression_ratio=wan_default_config[config.model_kwargs.model_name]["temporal_compression_ratio"], + target_fps=self.fps, + allow_padding=allow_padding, + min_latent_frames=min_latent_frames, + single_video_only=single_video_only, + independent_first_frame=getattr(config, "independent_first_frame", False), + return_image=getattr(config, "i2v", False), + max_chunks_per_shot=max_chunks_per_shot, + sample_warning_seconds=dataset_sample_warning_seconds, + sample_warning_interval_seconds=dataset_sample_warning_interval_seconds, + ) + if allow_padding and self.is_main_process: + print(f"[Padding] Variable-length training enabled: short videos will be padded with loss masking" + f" (min_latent_frames={min_latent_frames})") + if single_video_only and self.is_main_process: + print(f"[uniform_prompt] single_video_only enabled: each sample uses one video only") + # SP ranks in the same SP group need the same batch because they shard + # the sequence dimension. Use dp_rank for data parallel sampling. + random_seed = int(time.time()) % (2**31) * dist.get_rank() + if self.sequence_parallel_size > 1: + dp_rank = global_rank // self.sequence_parallel_size + sampler = torch.utils.data.distributed.DistributedSampler( + dataset, shuffle=True, drop_last=True, + rank=dp_rank, num_replicas=self.data_parallel_size, + seed=random_seed, + ) + else: + sampler = torch.utils.data.distributed.DistributedSampler( + dataset, shuffle=True, drop_last=True, + seed=random_seed, + ) + dataloader = torch.utils.data.DataLoader( + dataset, + batch_size=config.batch_size, + sampler=sampler, + num_workers=2, + prefetch_factor=1, + pin_memory=False, + persistent_workers=False, + collate_fn=multi_video_collate_fn, + ) + + # Eval dataloader: batch size defaults to 1 to keep validation memory predictable. + eval_data_path = getattr(config, "eval_data_path", config.data_path) + inference_num_frames = section_get(config, "evaluation", "num_frames", getattr(config, "inference_num_frames", 0)) + if isinstance(inference_num_frames, (list, tuple)): + inference_num_frames = inference_num_frames[0] if len(inference_num_frames) > 0 else 0 + eval_total_frames = ( + (inference_num_frames - 1) * wan_default_config[config.model_kwargs.model_name]["temporal_compression_ratio"] + 1 + if inference_num_frames > 0 else total_frames + ) + temporal_compression_ratio = wan_default_config[config.model_kwargs.model_name]["temporal_compression_ratio"] + first_chunk_frames = 1 + (num_frame_per_block - 1) * temporal_compression_ratio + subsequent_chunk_frames = num_frame_per_block * temporal_compression_ratio + num_blocks = 1 + (eval_total_frames - first_chunk_frames) // subsequent_chunk_frames + chunks_per_shot = getattr(config, "chunks_per_shot", 0) + scene_cut_prefix = getattr(config, "scene_cut_prefix", "The scene transitions. ") + if getattr(config, "i2v", False): + eval_dataset = MultiVideoConcatDataset( + data_dir=eval_data_path, + video_size=(frame_raw_height, frame_raw_width), + total_frames=eval_total_frames, + deterministic=True, + num_frame_per_block=num_frame_per_block, + temporal_compression_ratio=temporal_compression_ratio, + target_fps=self.fps, + allow_padding=allow_padding, + min_latent_frames=min_latent_frames, + single_video_only=single_video_only, + independent_first_frame=getattr(config, "independent_first_frame", False), + return_image=True, + max_chunks_per_shot=max_chunks_per_shot, + scene_cut_prefix=scene_cut_prefix, + sample_warning_seconds=dataset_sample_warning_seconds, + sample_warning_interval_seconds=dataset_sample_warning_interval_seconds, + ) + eval_collate = multi_video_collate_fn + else: + eval_dataset = MultiTextConcatDataset( + data_path=eval_data_path, + num_blocks=num_blocks, + chunks_per_shot=chunks_per_shot, + scene_cut_prefix=scene_cut_prefix, + deterministic=True, + ) + eval_collate = eval_collate_fn + if dist.get_rank() == 0: + print(f"Using {eval_dataset.__class__.__name__} for eval: {eval_data_path}, num_blocks={num_blocks}") + eval_sampler = torch.utils.data.distributed.DistributedSampler( + eval_dataset, shuffle=False, drop_last=False + ) + eval_dataloader = torch.utils.data.DataLoader( + eval_dataset, + batch_size=section_get(config, "evaluation", "val_batch_size", 1), + sampler=eval_sampler, + num_workers=0, + pin_memory=False, + persistent_workers=False, + collate_fn=eval_collate, + ) + + if dist.get_rank() == 0: + print("DATASET SIZE %d" % len(dataset)) + print("EVAL DATASET SIZE %d" % len(eval_dataset)) + + self.dataloader = cycle(dataloader) + self.eval_dataloader = eval_dataloader + + ############################################################################################################## + # 6. Set up EMA parameter containers + ema_weight = config.ema_weight + self.generator_ema = None + if (ema_weight is not None) and (ema_weight > 0.0) and (self.step >= config.ema_start_step): + if self.is_main_process: + print(f"Setting up EMA with weight {ema_weight}") + self.generator_ema = EMA_FSDP(self.model.generator, decay=ema_weight) + + ############################################################################################################## + # 7. (If resuming) Load optimizer and EMA from checkpoint + # Model weights were loaded before FSDP wrapping; restore only + # optimizer and EMA state that depend on FSDP here. + + if raw_state is not None: + if "generator_ema" in raw_state and self.generator_ema is not None: + self.generator_ema.load_state_dict(raw_state["generator_ema"]) + if self.is_main_process: + print("Resuming generator EMA...") + else: + if self.is_main_process: + print("Warning: Generator EMA checkpoint not found.") + + if "generator_optimizer" in raw_state: + gen_osd = FSDP.optim_state_dict_to_load( + self.model.generator, + self.generator_optimizer, + raw_state["generator_optimizer"], + ) + del raw_state["generator_optimizer"] + self.generator_optimizer.load_state_dict(gen_osd) + del gen_osd + if self.is_main_process: + print("Resuming generator optimizer...") + else: + if self.is_main_process: + print("Warning: Generator optimizer checkpoint not found.") + + del raw_state + gc.collect() + + ############################################################################################################## + + self.max_grad_norm = getattr(config, "max_grad_norm", 10.0) + self.previous_time = None + + # Resume error buffer from checkpoint. + # Try ``*_sp{sp_rank}.pt`` first, fall back to ``*.pt`` (legacy). + if self.model.error_buffer is not None and auto_resume: + ckpt_dir = self.find_latest_checkpoint(self.output_path) + if ckpt_dir is not None: + ckpt_root = os.path.dirname(ckpt_dir) + sp_size_ = max(self.sequence_parallel_size, 1) + global_rank = dist.get_rank() if dist.is_initialized() else 0 + sp_rank = global_rank % sp_size_ + + def _resolve_buf_file(stem): + if sp_size_ > 1: + p = os.path.join(ckpt_root, f"{stem}_sp{sp_rank}.pt") + if os.path.exists(p): + return p + p = os.path.join(ckpt_root, f"{stem}.pt") + return p if os.path.exists(p) else None + + for stem, buffer in [("error_buffer", self.model.error_buffer), + ("noise_error_buffer", self.model.noise_error_buffer)]: + if buffer is None: + continue + bf = _resolve_buf_file(stem) + if bf is not None: + bf_state = torch.load(bf, map_location="cpu") + buffer.load_state_dict(bf_state) + del bf_state + s = buffer.stats() + rng = s.get('global_block_range', '') + shard = s.get('shard', '') + print(f"[{stem}] rank={global_rank} Resumed from " + f"{os.path.basename(bf)}: {s['total_entries']} entries, " + f"{s['filled_buckets']} buckets, " + f"total_added={s['total_added']} {rng} {shard}".rstrip()) + elif self.is_main_process: + print(f"[{stem}] No saved buffer found, starting fresh.") + + def _move_optimizer_to_device(self, optimizer, device): + """Move optimizer state to the specified device.""" + for state in optimizer.state.values(): + for k, v in state.items(): + if isinstance(v, torch.Tensor): + state[k] = v.to(device) + + def find_latest_checkpoint(self, logdir): + """Find the latest checkpoint in the logdir.""" + if not os.path.exists(logdir): + return None + + checkpoint_dirs = [] + for item in os.listdir(logdir): + if item.startswith("checkpoint_model_") and os.path.isdir(os.path.join(logdir, item)): + try: + # Extract step number from directory name + step_str = item.replace("checkpoint_model_", "") + step = int(step_str) + checkpoint_path = os.path.join(logdir, item, "model.pt") + if os.path.exists(checkpoint_path): + checkpoint_dirs.append((step, checkpoint_path)) + except ValueError: + continue + + if not checkpoint_dirs: + return None + + # Sort by step number and return the latest one + checkpoint_dirs.sort(key=lambda x: x[0]) + latest_step, latest_path = checkpoint_dirs[-1] + return latest_path + + def get_all_checkpoints(self, logdir): + """Get all checkpoints in the logdir sorted by step number.""" + if not os.path.exists(logdir): + return [] + + checkpoint_dirs = [] + for item in os.listdir(logdir): + if item.startswith("checkpoint_model_") and os.path.isdir(os.path.join(logdir, item)): + try: + # Extract step number from directory name + step_str = item.replace("checkpoint_model_", "") + step = int(step_str) + checkpoint_dir_path = os.path.join(logdir, item) + checkpoint_file_path = os.path.join(checkpoint_dir_path, "model.pt") + if os.path.exists(checkpoint_file_path): + checkpoint_dirs.append((step, checkpoint_dir_path, item)) + except ValueError: + continue + + # Sort by step number (ascending order) + checkpoint_dirs.sort(key=lambda x: x[0]) + return checkpoint_dirs + + def cleanup_old_checkpoints(self, logdir, max_checkpoints): + """Remove old checkpoints if the number exceeds max_checkpoints. + + Only the main process performs the actual deletion to avoid race conditions + in distributed training. + """ + if max_checkpoints <= 0: + return + + # Only main process should perform cleanup to avoid race conditions + if not self.is_main_process: + return + + checkpoints = self.get_all_checkpoints(logdir) + if len(checkpoints) > max_checkpoints: + # Calculate how many to remove + num_to_remove = len(checkpoints) - max_checkpoints + checkpoints_to_remove = checkpoints[:num_to_remove] # Remove oldest ones + + print(f"Checkpoint cleanup: Found {len(checkpoints)} checkpoints, removing {num_to_remove} oldest ones (keeping {max_checkpoints})") + + import shutil + removed_count = 0 + for step, checkpoint_dir_path, dir_name in checkpoints_to_remove: + try: + print(f" Removing: {dir_name} (step {step})") + shutil.rmtree(checkpoint_dir_path) + removed_count += 1 + except Exception as e: + print(f" Warning: Failed to remove checkpoint {dir_name}: {e}") + + print(f"Checkpoint cleanup completed: removed {removed_count}/{num_to_remove} old checkpoints") + else: + if len(checkpoints) > 0: + print(f"Checkpoint cleanup: Found {len(checkpoints)} checkpoints (max: {max_checkpoints}, no cleanup needed)") + + def save(self): + print("Start gathering distributed model states...") + + # Release large inference caches before saving when possible. + if hasattr(self.model, "inference_pipeline") and self.model.inference_pipeline is not None: + clear_fn = getattr(self.model.inference_pipeline, "clear_cache", None) + if clear_fn is not None: + try: + clear_fn() + except Exception as e: + print(f"Warning: failed to clear inference cache before save: {e}") + # Drop the inference pipeline reference so GC / empty_cache can + # reclaim memory. + self.model.inference_pipeline = None + torch.cuda.empty_cache() + + with FSDP.state_dict_type( + self.model.generator, + StateDictType.FULL_STATE_DICT, + FullStateDictConfig(rank0_only=True, offload_to_cpu=True), + FullOptimStateDictConfig(rank0_only=True, offload_to_cpu=True), + ): + generator_state_dict = self.model.generator.state_dict() + generator_opim_state_dict = FSDP.optim_state_dict(self.model.generator, + self.generator_optimizer) + + if self.config.ema_start_step < self.step and self.generator_ema is not None: + state_dict = { + "generator": generator_state_dict, + "generator_ema": self.generator_ema.state_dict(), + "generator_optimizer": generator_opim_state_dict, + "step": self.step, + } + else: + state_dict = { + "generator": generator_state_dict, + "generator_optimizer": generator_opim_state_dict, + "step": self.step, + } + + checkpoint_dir = os.path.join(self.output_path, f"checkpoint_model_{self.step:06d}") + if self.is_main_process: + os.makedirs(checkpoint_dir, exist_ok=True) + checkpoint_file = os.path.join(checkpoint_dir, "model.pt") + torch.save(state_dict, checkpoint_file) + print("Model saved to", checkpoint_file) + + # Save error buffer — unified per-sp_rank pattern: + # Each SP rank owns a different t-bucket shard (and different + # positions in 2D mode). The first DP rank in each SP group + # writes ``error_buffer_sp{sp_rank}.pt``. + # Fallback (sp_size<=1): main_process writes ``error_buffer.pt``. + if self.model.error_buffer is not None: + sp_size_ = max(self.sequence_parallel_size, 1) + _global_rank = dist.get_rank() if dist.is_initialized() else 0 + _sp_rank = _global_rank % sp_size_ + _is_first_dp = (_global_rank // sp_size_) == 0 + + if dist.is_initialized(): + dist.barrier() + + should_save = _is_first_dp if sp_size_ > 1 else self.is_main_process + if should_save: + for stem, buffer in [("error_buffer", self.model.error_buffer), + ("noise_error_buffer", self.model.noise_error_buffer)]: + if buffer is None: + continue + fname = f"{stem}_sp{_sp_rank}.pt" if sp_size_ > 1 else f"{stem}.pt" + fpath = os.path.join(checkpoint_dir, fname) + torch.save(buffer.state_dict(), fpath) + s = buffer.stats() + rng = s.get('global_block_range', '') + shard = s.get('shard', '') + print(f"[rank={_global_rank}] {stem} saved to {fname} " + f"({s['total_entries']} entries, {s['filled_buckets']} buckets) " + f"{rng} {shard}".rstrip()) + + if self.is_main_process: + # Cleanup old checkpoints if max_checkpoints is set + max_checkpoints = getattr(self.config, "max_checkpoints", 0) + if max_checkpoints > 0: + self.cleanup_old_checkpoints(self.output_path, max_checkpoints) + + # Keep all ranks in sync so non-rank0 workers don't kick off the next + # training iteration (and trigger NCCL watchdog timeouts) while rank0 + # is still writing the checkpoint to disk. + if dist.is_initialized(): + dist.barrier() + + torch.cuda.empty_cache() + import gc + gc.collect() + + def train_one_step(self, batch, accumulation_step=0, accumulation_steps=None): + accumulation_steps = accumulation_steps or getattr(self, "gradient_accumulation_steps", 1) + self.log_iters = 1 + + if self.step % 20 == 0: + torch.cuda.empty_cache() + # Step 1: Get the next batch of text prompts + text_prompts = batch["prompts"] + batch_size = len(text_prompts) + clean_latent_is_sp_sharded = False + if not self.config.load_raw_video: # precomputed latent + clean_latent = batch["ode_latent"][:, -1].to( + device=self.device, dtype=self.dtype) + image_latent = clean_latent[:, 0:1] + else: # encode raw video to latent + ( + clean_latent, + image_latent, + clean_latent_is_sp_sharded, + ) = self.sp_helper.encode_raw_video_latents( + batch, + batch_size=batch_size, + ) + + loss_mask = self.sp_helper.build_loss_mask( + batch, clean_latent, clean_latent_is_sp_sharded + ) + image_or_video_shape = list(self.config.image_or_video_shape) + image_or_video_shape[0] = batch_size + # Step 2: Extract the conditional infos + with torch.no_grad(): + # turn text prompts: List[List[str]] into List[str] + text_prompts_flat = [prompt for sublist in text_prompts for prompt in sublist] + + conditional_dict = self.model.text_encoder( + text_prompts=text_prompts_flat) + + if not getattr(self, "unconditional_dict", None): + unconditional_dict = self.model.text_encoder( + text_prompts=[self.config.negative_prompt] * batch_size) + unconditional_dict = {k: v.detach() + for k, v in unconditional_dict.items()} + self.unconditional_dict = unconditional_dict # cache the unconditional_dict + else: + unconditional_dict = self.unconditional_dict + + # Step 2.5: Sequence Parallel partitions sequence-owned tensors. + if self.sequence_parallel_size > 1: + clean_latent, conditional_dict, image_or_video_shape = ( + self.sp_helper.partition_training_inputs( + image_or_video_shape=image_or_video_shape, + clean_latent=clean_latent, + conditional_dict=conditional_dict, + clean_latent_is_sharded=clean_latent_is_sp_sharded, + ) + ) + image_latent = self.sp_helper.local_i2v_initial_latent(image_latent) + loss_mask, loss_mask_global_valid_count = self.sp_helper.partition_loss_mask( + loss_mask, + already_sharded=clean_latent_is_sp_sharded, + ) + + # Step 3: Train the generator + gen_kwargs = dict( + image_or_video_shape=image_or_video_shape, + conditional_dict=conditional_dict, + unconditional_dict=unconditional_dict, + clean_latent=clean_latent, + initial_latent=image_latent, + loss_mask=loss_mask, + loss_mask_global_valid_count=loss_mask_global_valid_count, + global_step=self.step, + ) + generator_loss, log_dict = self.model.generator_loss(**gen_kwargs) + if accumulation_step == 0: + self.generator_optimizer.zero_grad(set_to_none=True) + scaled_loss = generator_loss / accumulation_steps + scaled_loss.backward() + if accumulation_step == accumulation_steps - 1: + generator_grad_norm = self.model.generator.clip_grad_norm_( + self.max_grad_norm) + + self.generator_optimizer.step() + self.step += 1 + else: + generator_grad_norm = torch.tensor(0.0, device=self.device) + + # Run the remaining logic only after a full gradient-accumulation cycle. + if accumulation_step != accumulation_steps - 1: + return + + # Step 4: Update EMA (if enabled and after start step) + if (self.step >= self.config.ema_start_step) and \ + (self.generator_ema is None) and \ + (getattr(self.config, "ema_weight", None) is not None) and \ + (self.config.ema_weight > 0): + self.generator_ema = EMA_FSDP(self.model.generator, decay=self.config.ema_weight) + + # Update EMA after optimizer step + if self.generator_ema is not None and self.step >= self.config.ema_start_step: + self.generator_ema.update(self.model.generator) + + wandb_loss_dict = { + "generator_loss": generator_loss.item(), + "generator_grad_norm": generator_grad_norm.item(), + } + + # Error buffer stats + er_log_str = "" + if "er_total_added" in log_dict: + wandb_loss_dict["er_total_entries"] = log_dict["er_total_entries"] + wandb_loss_dict["er_total_added"] = log_dict["er_total_added"] + wandb_loss_dict["er_injected"] = int(log_dict["er_injected"]) + wandb_loss_dict["er_latent_injected"] = int(log_dict["er_latent_injected"]) + wandb_loss_dict["er_noise_injected"] = int(log_dict.get("er_noise_injected", False)) + wandb_loss_dict["er_noise_total_entries"] = log_dict.get("er_noise_total_entries", 0) + ctx_flag = 'Y' if log_dict['er_injected'] else 'N' + lat_flag = 'Y' if log_dict['er_latent_injected'] else 'N' + noise_flag = 'Y' if log_dict.get('er_noise_injected', False) else 'N' + er_log_str = ( + f", er_buf={log_dict['er_total_entries']}|" + f"{log_dict.get('er_noise_total_entries', 0)} " + f"({log_dict['er_filled_buckets']} buckets), " + f"ctx={ctx_flag} lat={lat_flag} noise={noise_flag}" + ) + + # Step 5: Logging + if self.is_main_process: + if not self.disable_wandb: + wandb.log(wandb_loss_dict, step=self.step) + print( + f"[step {self.step:07d}] " + f"generator_loss={wandb_loss_dict['generator_loss']:.6f}, " + f"generator_grad_norm={wandb_loss_dict['generator_grad_norm']:.6f}" + f"{er_log_str}" + ) + + if self.step % self.config.gc_interval == 0: + if dist.get_rank() == 0: + logging.info("DistGarbageCollector: Running GC.") + gc.collect() + + def _set_sp_attn(self, enabled: bool): + """ + Toggle SP self-attention between training and inference. + This only applies to 5B runs with SP enabled. + """ + if not hasattr(self, "_sp_attn_blocks"): + return + if self.sequence_parallel_size <= 1: + return + + # Lazy import to avoid failures under non-5B configurations. + try: + from wan_5b.distributed.sequence_parallel import sp_causal_attn_forward + except Exception: + return + + for sa in self._sp_attn_blocks: + if not hasattr(sa, "_orig_forward"): + continue + if enabled: + sa.forward = types.MethodType(sp_causal_attn_forward, sa) + else: + sa.forward = sa._orig_forward + + @torch.no_grad() + def _swap_ema_weights(self): + """ + Bidirectionally swap model weights with EMA shadow weights. + Calling this twice restores both the model and EMA to their original state. + """ + with FSDP.summon_full_params(self.model.generator, writeback=True): + for n, p in self.model.generator.module.named_parameters(): + cleaned_name = EMA_FSDP._clean_param_name(n) + if cleaned_name in self.generator_ema.shadow: + ema_val = self.generator_ema.shadow[cleaned_name] + tmp = p.data.clone().float().cpu() + p.data.copy_(ema_val.to(dtype=p.dtype, device=p.device)) + self.generator_ema.shadow[cleaned_name] = tmp + + def _run_evaluation_inference(self): + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + + if self.model.inference_pipeline is None: + self.model._initialize_inference_pipeline() + + out_dir = os.path.join(self.output_path, f"generated_video_{self.step:06d}") + if self.is_main_process: + os.makedirs(out_dir, exist_ok=True) + barrier() + + rank = dist.get_rank() + vis_ema = section_get(self.config, "evaluation", "use_ema", getattr(self.config, "vis_ema", False)) + vis_ema = vis_ema and self.generator_ema is not None + + for eval_batch in self.eval_dataloader: + eval_prompts = eval_batch["prompts"] + eval_idx = eval_batch["idx"] + eval_images = eval_batch.get("image", None) + + batch_size_eval = len(eval_prompts) + for b in range(batch_size_eval): + prompts_for_sample = eval_prompts[b] + + if self.is_main_process: + print(f"prompts_for_sample: {prompts_for_sample}") + print(len(prompts_for_sample)) + print(prompts_for_sample[0][:60]) + + sample_idx = ( + eval_idx[b].item() + if hasattr(eval_idx, "shape") + else int(eval_idx[b]) + ) + + save_latents_only = section_get( + self.config, + "evaluation", + "save_latents_only", + self.config.get("return_latents", False), + aliases=("return_latents", "save_latent_only"), + ) + + run_modes = [("", False)] + if vis_ema: + run_modes.append(("_ema", True)) + + for suffix, use_ema in run_modes: + generated_video = self.generate_video( + self.model.inference_pipeline, + [prompts_for_sample], + eval_images[b:b + 1] if eval_images is not None else None, + use_ema=use_ema, + ) + + if not save_latents_only: + video_path = os.path.join( + out_dir, + f"video{suffix}_rank{rank:02d}_idx{sample_idx:06d}.mp4", + ) + write_video(video_path, generated_video[0], fps=self.fps) + else: + video_path = os.path.join( + out_dir, + f"latents{suffix}_rank{rank:02d}_idx{sample_idx:06d}.pt", + ) + torch.save(generated_video[0], video_path) + + if (not self.disable_wandb) and self.is_main_process and not save_latents_only: + caption = prompts_for_sample[0] if len(prompts_for_sample) > 0 else "" + log_key = f"generated_video{suffix}" + wandb.log( + { + log_key: wandb.Video( + generated_video[0].transpose(0, 3, 1, 2), + caption=f"{caption}", + fps=self.fps, + format="mp4", + ), + }, + step=self.step, + ) + + del generated_video + + prompt_txt_path = os.path.join( + out_dir, + f"prompt_rank{rank:02d}_idx{sample_idx:06d}.txt", + ) + save_prompts_to_txt( + prompts_for_sample, + prompt_txt_path, + self.is_main_process, + ) + barrier() + + if hasattr(self.model, "inference_pipeline") and self.model.inference_pipeline is not None: + clear_fn = getattr(self.model.inference_pipeline, "clear_cache", None) + if clear_fn is not None: + clear_fn() + torch.cuda.empty_cache() + + @torch.no_grad() + def generate_video(self, pipeline, prompts, image=None, use_ema=False): + # Temporarily disable SP self-attention during inference to avoid + # interfering with KV-cache logic. + self._set_sp_attn(False) + ema_applied = use_ema and self.generator_ema is not None + if ema_applied: + self._swap_ema_weights() + try: + batch_size = len(prompts) + noise_shape = list(self.config.image_or_video_shape[1:]) + inference_num_frames = section_get( + self.config, "evaluation", "num_frames", getattr(self.config, "inference_num_frames", 0) + ) + if isinstance(inference_num_frames, (list, tuple)): + inference_num_frames = inference_num_frames[0] if len(inference_num_frames) > 0 else 0 + if inference_num_frames > 0: + noise_shape[0] = inference_num_frames + initial_latent = None + if image is not None: + image = image.to(device="cuda", dtype=self.dtype) + if image.ndim == 4: + image = image.unsqueeze(2) + elif image.ndim != 5: + raise ValueError(f"Expected i2v image with shape [B,C,H,W] or [B,C,T,H,W], got {tuple(image.shape)}") + initial_latent = pipeline.vae.encode_to_latent(image).to(device="cuda", dtype=self.dtype) + if initial_latent.shape[0] != batch_size: + initial_latent = initial_latent.repeat(batch_size, 1, 1, 1, 1) + if noise_shape[0] <= initial_latent.shape[1]: + raise ValueError( + f"evaluation.num_frames must exceed the i2v conditioning frames; " + f"got {inference_num_frames} and {initial_latent.shape[1]}" + ) + sampled_noise = torch.randn( + [batch_size] + noise_shape, device="cuda", dtype=self.dtype + ) + + save_latents_only = section_get( + self.config, + "evaluation", + "save_latents_only", + self.config.get("return_latents", False), + aliases=("return_latents", "save_latent_only"), + ) + video = pipeline.inference( + noise=sampled_noise, + text_prompts=prompts, + initial_latent=initial_latent, + return_latents=save_latents_only + ) + if not save_latents_only: + current_video = video.permute(0, 1, 3, 4, 2).cpu().numpy() * 255.0 + else: + current_video = video + finally: + if ema_applied: + self._swap_ema_weights() + # Restore SP self-attention for training. + self._set_sp_attn(True) + + return current_video + + def _sync_batch_for_sequence_parallel(self, batch, accumulation_step: int = 0): + return self.sp_helper.sync_batch(batch, step=self.step) + + def train(self): + if getattr(self.config, "generate_before_train", False): + if self.is_main_process: + print("[generate_before_train] Running evaluation inference before training starts...") + self._run_evaluation_inference() + if self.is_main_process: + print("[generate_before_train] Inference done. Exiting.") + barrier() + return + + acc_steps = getattr(self, "gradient_accumulation_steps", 1) + while True: + for acc in range(acc_steps): + batch = next(self.dataloader) + + # Synchronize batch contents across ranks under Sequence Parallel. + if self.sequence_parallel_size > 1: + batch = self._sync_batch_for_sequence_parallel(batch, accumulation_step=acc) + + self.train_one_step(batch, accumulation_step=acc, accumulation_steps=acc_steps) + if (not self.config.no_save) and self.step % self.config.log_iters == 0: + torch.cuda.empty_cache() + self.save() + torch.cuda.empty_cache() + + evaluation_interval = section_get(self.config, "evaluation", "interval", getattr(self.config, "generate_interval", 0)) + if evaluation_interval > 0 and self.step % evaluation_interval == 0: + self._run_evaluation_inference() + + barrier() + if self.is_main_process: + current_time = time.time() + if self.previous_time is None: + self.previous_time = current_time + else: + if not self.disable_wandb: + wandb.log({"per iteration time": current_time - self.previous_time}, step=self.step) + self.previous_time = current_time + + if self.step >= self.config.max_iters: + break diff --git a/trainer/distillation.py b/trainer/distillation.py new file mode 100644 index 0000000..d61ce9a --- /dev/null +++ b/trainer/distillation.py @@ -0,0 +1,1483 @@ +# Adopted from https://github.com/guandeh17/Self-Forcing +# SPDX-License-Identifier: Apache-2.0 +import gc +import logging + +from utils.dataset import cycle +from utils.dataset import MultiVideoConcatDataset, MultiTextConcatDataset, multi_video_collate_fn, eval_collate_fn, DEFAULT_SCENE_CUT_PREFIX +from utils.config import section_get, wan_default_config +from utils.distributed import EMA_FSDP, fsdp_wrap, launch_distributed_job +from utils.misc import ( + set_seed, + merge_dict_list +) +import torch.distributed as dist +from omegaconf import OmegaConf +from model import DMD +import torch +import wandb +import os +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +from torch.distributed.fsdp import ( + StateDictType, FullStateDictConfig, FullOptimStateDictConfig +) +from torchvision.io import write_video + +# LoRA related imports +import peft +from peft import get_peft_model_state_dict + +from pipeline import ( + CausalDiffusionInferencePipeline +) +import time + +class Trainer: + + def __init__(self, config): + self.config = config + self.step = 0 + + # Step 1: Initialize the distributed training environment (rank, seed, dtype, logging etc.) + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + + launch_distributed_job() + global_rank = dist.get_rank() + self.world_size = dist.get_world_size() + + self.dtype = torch.bfloat16 if config.mixed_precision else torch.float32 + self.device = torch.cuda.current_device() + self.is_main_process = global_rank == 0 + self.causal = getattr(config, "causal", getattr(config, "all_causal", True)) + self.disable_wandb = config.disable_wandb + + # use a random seed for the training + if config.seed == 0: + random_seed = torch.randint(0, 10000000, (1,), device=self.device) + dist.broadcast(random_seed, src=0) + config.seed = random_seed.item() + + set_seed(config.seed + global_rank) + + if self.is_main_process and not self.disable_wandb: + if getattr(config, "wandb_key", None): + wandb.login(key=config.wandb_key) + wandb.init( + config=OmegaConf.to_container(config, resolve=True), + name=config.config_name, + id=config.config_name, + mode="online", + entity=config.wandb_entity, + project=config.wandb_project, + dir=config.wandb_save_dir, + resume="allow" + ) + + self.output_path = config.logdir + + # Step 2: Initialize the model + if config.distribution_loss == "dmd": + self.model = DMD(config, device=self.device) + else: + raise ValueError(f"Unsupported distribution matching loss: {config.distribution_loss}") + + # Save pretrained model state_dicts to CPU + self.fake_score_state_dict_cpu = self.model.fake_score.state_dict() + + # ================================= NVFP4 Quantized Training / Inference ================================= + # `generator_quant` is the preferred student flag; `model_quant` is kept + # as a legacy alias used by earlier Sage configs. + self.generator_quant = getattr(config, "generator_quant", getattr(config, "model_quant", False)) + self.real_score_quant = getattr(config, "real_score_quant", False) + self.fake_score_quant = getattr(config, "fake_score_quant", False) + self.real_score_quant_materialize = getattr(config, "real_score_quant_materialize", True) + + if self.generator_quant or self.real_score_quant or self.fake_score_quant: + from utils.quant import ModelQuantizationConfig, quantize_model_with_filter + + fallback_sr = getattr(config, "model_quant_scale_rule", "static_6") + fallback_asr = getattr(config, "model_quant_activation_scale_rule", "static_6") + fallback_wsr = getattr(config, "model_quant_weight_scale_rule", None) + fallback_gsr = getattr(config, "model_quant_gradient_scale_rule", None) + + if self.generator_quant: + gen_cfg = ModelQuantizationConfig( + scale_rule=getattr(config, "generator_quant_scale_rule", fallback_sr), + activation_scale_rule=getattr(config, "generator_quant_activation_scale_rule", fallback_asr), + weight_scale_rule=getattr(config, "generator_quant_weight_scale_rule", fallback_wsr), + gradient_scale_rule=getattr(config, "generator_quant_gradient_scale_rule", fallback_gsr), + keep_master_weights=True, + weight_scale_2d=True, + ) + self.model.generator.model, gen_matched = quantize_model_with_filter( + self.model.generator.model, + quant_config=gen_cfg, + filtered_modules=getattr(config, "generator_quant_filtered_modules", None), + filter_profile="student", + use_default_filtered_modules=getattr( + config, "generator_quant_use_default_filtered_modules", True + ), + cast_model_to_bf16=False, + materialize_for_inference=False, + verbose=self.is_main_process, + ) + if self.is_main_process: + print( + "[NVFP4] Generator (student) quantized training enabled, " + f"scale_rule={gen_cfg.scale_rule}, {len(gen_matched)} modules excluded" + ) + + if self.real_score_quant: + real_cfg = ModelQuantizationConfig( + scale_rule=getattr(config, "real_score_quant_scale_rule", fallback_sr), + activation_scale_rule=getattr(config, "real_score_quant_activation_scale_rule", fallback_asr), + weight_scale_rule=getattr(config, "real_score_quant_weight_scale_rule", fallback_wsr), + gradient_scale_rule=None, + keep_master_weights=True, + weight_scale_2d=True, + ) + self.model.real_score.model, real_matched = quantize_model_with_filter( + self.model.real_score.model, + quant_config=real_cfg, + filtered_modules=getattr(config, "real_score_quant_filtered_modules", None), + filter_profile="teacher", + use_default_filtered_modules=getattr( + config, "real_score_quant_use_default_filtered_modules", True + ), + cast_model_to_bf16=False, + materialize_for_inference=False, + verbose=self.is_main_process, + ) + if self.is_main_process: + real_score_plan = ( + "auto materialize before FSDP wrapping" + if self.real_score_quant_materialize + else "keep master weights after checkpoint load" + ) + print( + "[NVFP4] Real_score (teacher) quantized inference enabled, " + f"scale_rule={real_cfg.scale_rule}, {len(real_matched)} modules excluded, " + f"{real_score_plan}" + ) + + if self.fake_score_quant: + fake_cfg = ModelQuantizationConfig( + scale_rule=getattr(config, "fake_score_quant_scale_rule", fallback_sr), + activation_scale_rule=getattr(config, "fake_score_quant_activation_scale_rule", fallback_asr), + weight_scale_rule=getattr(config, "fake_score_quant_weight_scale_rule", fallback_wsr), + gradient_scale_rule=getattr(config, "fake_score_quant_gradient_scale_rule", fallback_gsr), + keep_master_weights=True, + weight_scale_2d=True, + ) + self.model.fake_score.model, fake_matched = quantize_model_with_filter( + self.model.fake_score.model, + quant_config=fake_cfg, + filtered_modules=getattr(config, "fake_score_quant_filtered_modules", None), + filter_profile="critic", + use_default_filtered_modules=getattr( + config, "fake_score_quant_use_default_filtered_modules", True + ), + cast_model_to_bf16=False, + materialize_for_inference=False, + verbose=self.is_main_process, + ) + if self.is_main_process: + print( + "[NVFP4] Fake_score (critic) quantized training enabled, " + f"scale_rule={fake_cfg.scale_rule}, {len(fake_matched)} modules excluded" + ) + + # Auto resume configuration (needed for LoRA checkpoint loading) + auto_resume = getattr(config, "auto_resume", True) # Default to True + + # ================================= LoRA Configuration ================================= + self.is_lora_enabled = False + self.lora_config = None + + if hasattr(config, 'adapter') and config.adapter is not None: + self.is_lora_enabled = True + self.lora_config = config.adapter + + if self.is_main_process: + print(f"LoRA enabled with config: {self.lora_config}") + print("Loading base model and applying LoRA before FSDP wrapping...") + + # 1. Load base model first (config.generator_ckpt) before applying LoRA. + generator_checkpoint_path = getattr(config, "generator_ckpt", None) + if generator_checkpoint_path: + if self.is_main_process: + print(f"Loading base model from {generator_checkpoint_path} (before applying LoRA)") + generator_checkpoint = torch.load(generator_checkpoint_path, map_location="cpu") + + # Load generator (directly; no key alignment needed since LoRA not applied yet) + if isinstance(generator_checkpoint, dict) and "generator" in generator_checkpoint: + if self.is_main_process: + print(f"Loading pretrained generator from {generator_checkpoint_path}") + self.model.generator.load_state_dict(generator_checkpoint["generator"], strict=True) + if self.is_main_process: + print("Generator weights loaded successfully") + elif isinstance(generator_checkpoint, dict) and "model" in generator_checkpoint: + if self.is_main_process: + print(f"Loading pretrained generator from {generator_checkpoint_path}") + self.model.generator.load_state_dict(generator_checkpoint["model"], strict=True) + if self.is_main_process: + print("Generator weights loaded successfully") + else: + self.model.generator.load_state_dict(generator_checkpoint, strict=True) + if self.is_main_process: + print("Loading base model as raw state_dict") + + # Load critic from full/base checkpoints when available. + if isinstance(generator_checkpoint, dict) and "critic" in generator_checkpoint: + if self.is_main_process: + print(f"Loading pretrained critic from {generator_checkpoint_path}") + self.model.fake_score.load_state_dict(generator_checkpoint["critic"], strict=True) + if self.is_main_process: + print("Critic weights loaded successfully") + # Load training step from checkpoint metadata. + if isinstance(generator_checkpoint, dict) and "step" in generator_checkpoint: + self.step = generator_checkpoint["step"] + if self.is_main_process: + print(f"base_checkpoint step: {self.step}") + else: + if self.is_main_process: + print("Warning: Step not found in checkpoint, starting from step 0.") + del generator_checkpoint + gc.collect() + else: + if self.is_main_process: + print("No base model checkpoint specified, skipping base weight loading for LoRA training.") + + # Load real_score from a separate checkpoint (independent of LoRA / auto_resume) + real_score_ckpt = getattr(config, "real_score_ckpt", None) + if real_score_ckpt: + if self.is_main_process: + print(f"Loading real_score from {real_score_ckpt}") + real_ckpt = torch.load(real_score_ckpt, map_location="cpu") + if "generator" in real_ckpt: + self.model.real_score.load_state_dict(real_ckpt["generator"], strict=True) + elif "critic" in real_ckpt: + self.model.real_score.load_state_dict(real_ckpt["critic"], strict=True) + elif "model" in real_ckpt: + self.model.real_score.load_state_dict(real_ckpt["model"], strict=True) + else: + if self.is_main_process: + print(f"No recognized key in {real_score_ckpt}, treating as raw state_dict") + self.model.real_score.load_state_dict(real_ckpt, strict=True) + del real_ckpt + gc.collect() + if self.is_main_process: + print(f"Successfully loaded real_score from {real_score_ckpt}") + + # Apply LoRA wrapping if enabled (after all base weights are loaded, before FSDP) + if self.is_lora_enabled: + # 2. Apply LoRA wrapping now (after loading base model, before FSDP wrapping) + if self.is_main_process: + print("Applying LoRA to models...") + self.model.generator.model = self._configure_lora_for_model(self.model.generator.model, "generator") + + # Configure LoRA for fake_score if needed + if getattr(self.lora_config, 'apply_to_critic', True): + self.model.fake_score.model = self._configure_lora_for_model(self.model.fake_score.model, "fake_score") + if self.is_main_process: + print("LoRA applied to both generator and critic") + else: + if self.is_main_process: + print("LoRA applied to generator only") + + # 3. Load LoRA weights before FSDP wrapping (if a checkpoint is available). + # Priority: auto_resume -> legacy lora_ckpt -> initialized adapters. + lora_checkpoint_path = None + lora_checkpoint = None + if auto_resume and self.output_path: + latest_checkpoint = self.find_latest_checkpoint(self.output_path) + if latest_checkpoint: + lora_checkpoint_path = latest_checkpoint + if self.is_main_process: + print(f"Auto resume: Found LoRA checkpoint at {lora_checkpoint_path}") + else: + if self.is_main_process: + print("Auto resume: No LoRA checkpoint found in logdir") + elif auto_resume: + if self.is_main_process: + print("Auto resume enabled but no logdir specified for LoRA") + else: + if self.is_main_process: + print("Auto resume disabled for LoRA") + + if lora_checkpoint_path is not None: + lora_checkpoint = torch.load(lora_checkpoint_path, map_location="cpu") + elif getattr(config, "lora_ckpt", None): + lora_checkpoint_path = config.lora_ckpt + lora_checkpoint = torch.load(lora_checkpoint_path, map_location="cpu") + if self.is_main_process: + print(f"Using legacy lora_ckpt: {lora_checkpoint_path}") + elif self.is_main_process: + print("No LoRA checkpoint specified, starting LoRA training from scratch") + + # Load LoRA checkpoint (before FSDP wrapping) + if lora_checkpoint is not None: + if self.is_main_process: + print(f"Loading LoRA checkpoint from {lora_checkpoint_path} (before FSDP wrapping)") + + if "generator_lora" not in lora_checkpoint: + raise ValueError(f"LoRA checkpoint {lora_checkpoint_path} is not a valid LoRA checkpoint. " + f"Found keys: {list(lora_checkpoint.keys())}") + + if self.is_main_process: + print(f"Loading LoRA generator weights: {len(lora_checkpoint['generator_lora'])} keys in checkpoint") + peft.set_peft_model_state_dict(self.model.generator.model, lora_checkpoint["generator_lora"]) + del lora_checkpoint["generator_lora"] + + if getattr(self.lora_config, 'apply_to_critic', True): + if "critic_lora" not in lora_checkpoint: + raise ValueError(f"LoRA checkpoint {lora_checkpoint_path} is missing critic_lora.") + if self.is_main_process: + print(f"Loading LoRA critic weights: {len(lora_checkpoint['critic_lora'])} keys in checkpoint") + peft.set_peft_model_state_dict(self.model.fake_score.model, lora_checkpoint["critic_lora"]) + del lora_checkpoint["critic_lora"] + gc.collect() + + if "step" in lora_checkpoint: + self.step = lora_checkpoint["step"] + if self.is_main_process: + print(f"Resuming LoRA training from step {self.step}") + else: + if self.is_main_process: + print("No LoRA checkpoint to load, starting from scratch") + + # Materialize quantized inference-only weights before FSDP can expose + # sharded 1D parameter views. The student/critic are materialized only + # in LoRA mode because their base weights are frozen after adapters. + if self.generator_quant and self.is_lora_enabled: + self._materialize_quantized_model_before_fsdp( + self.model.generator.model, + "Generator", + cache_transposed_weights=True, + ) + + apply_lora_to_critic = getattr(self.lora_config, "apply_to_critic", True) if self.lora_config else False + if self.fake_score_quant and self.is_lora_enabled and apply_lora_to_critic: + self._materialize_quantized_model_before_fsdp( + self.model.fake_score.model, + "Fake_score", + cache_transposed_weights=True, + ) + + if self.real_score_quant and self.real_score_quant_materialize: + self._materialize_quantized_model_before_fsdp( + self.model.real_score.model, + "Real_score", + cache_transposed_weights=False, + ) + + self.model.generator = fsdp_wrap( + self.model.generator, + sharding_strategy=config.sharding_strategy, + mixed_precision=config.mixed_precision, + wrap_strategy=config.generator_fsdp_wrap_strategy + ) + + self.model.real_score = fsdp_wrap( + self.model.real_score, + sharding_strategy=config.sharding_strategy, + mixed_precision=config.mixed_precision, + wrap_strategy=config.real_score_fsdp_wrap_strategy + ) + + self.model.fake_score = fsdp_wrap( + self.model.fake_score, + sharding_strategy=config.sharding_strategy, + mixed_precision=config.mixed_precision, + wrap_strategy=config.fake_score_fsdp_wrap_strategy + ) + + self.model.text_encoder = fsdp_wrap( + self.model.text_encoder, + sharding_strategy=config.sharding_strategy, + mixed_precision=config.mixed_precision, + wrap_strategy=config.text_encoder_fsdp_wrap_strategy, + cpu_offload=getattr(config, "text_encoder_cpu_offload", False) + ) + self.model.vae = self.model.vae.to( + device=self.device, dtype=torch.bfloat16 if config.mixed_precision else torch.float32) + + # Step 3: Set up EMA parameter containers + rename_param = ( + lambda name: name.replace("_fsdp_wrapped_module.", "") + .replace("_checkpoint_wrapped_module.", "") + .replace("_orig_mod.", "") + ) + self.name_to_trainable_params = {} + for n, p in self.model.generator.named_parameters(): + if not p.requires_grad: + continue + + renamed_n = rename_param(n) + self.name_to_trainable_params[renamed_n] = p + ema_weight = config.ema_weight + self.generator_ema = None + if (ema_weight is not None) and (ema_weight > 0.0): + if self.is_lora_enabled: + if self.is_main_process: + print(f"EMA disabled in LoRA mode (LoRA provides efficient parameter updates without EMA)") + self.generator_ema = None + else: + print(f"Setting up EMA with weight {ema_weight}") + self.generator_ema = EMA_FSDP(self.model.generator, decay=ema_weight) + + # Step 4: Initialize the optimizer + self.generator_optimizer = torch.optim.AdamW( + [param for param in self.model.generator.parameters() + if param.requires_grad], + lr=config.lr, + betas=(config.beta1, config.beta2), + weight_decay=config.weight_decay + ) + + self.critic_optimizer = torch.optim.AdamW( + [param for param in self.model.fake_score.parameters() + if param.requires_grad], + lr=config.lr_critic if hasattr(config, "lr_critic") else config.lr, + betas=(config.beta1_critic, config.beta2_critic), + weight_decay=config.weight_decay + ) + + # Step 5: Initialize the dataloader + self.use_backward_simulation = getattr(config, "backward_simulation", True) + + model_name = config.model_kwargs.model_name + frame_raw_height = list(config.image_or_video_shape)[3] * wan_default_config[model_name]["spatial_compression_ratio"] + frame_raw_width = list(config.image_or_video_shape)[4] * wan_default_config[model_name]["spatial_compression_ratio"] + num_frame_per_block = getattr(config, "num_frame_per_block", 1) + self.fps = wan_default_config[model_name].get("fps", 16) + + latent_frames_for_dataset = list(config.image_or_video_shape)[1] + num_training_frames = getattr(config, "num_training_frames", latent_frames_for_dataset) + assert latent_frames_for_dataset >= num_training_frames, ( + f"image_or_video_shape[1] ({latent_frames_for_dataset}) must be >= " + f"num_training_frames ({num_training_frames}), otherwise the dataset " + f"will not provide enough prompts for the rollout." + ) + total_frames = (latent_frames_for_dataset - 1) * wan_default_config[model_name]["temporal_compression_ratio"] + 1 + if dist.get_rank() == 0: + print(f"[Dataset] latent_frames_for_dataset={latent_frames_for_dataset}, total_frames={total_frames}") + + temporal_compression_ratio = wan_default_config[model_name]["temporal_compression_ratio"] + first_chunk_frames = 1 + (num_frame_per_block - 1) * temporal_compression_ratio + subsequent_chunk_frames = num_frame_per_block * temporal_compression_ratio + num_blocks = 1 + (total_frames - first_chunk_frames) // subsequent_chunk_frames + if not getattr(config, "generator_is_causal", True): + num_blocks = 1 + + chunks_per_shot = getattr(config, "chunks_per_shot", 0) + scene_cut_prefix = getattr(config, "scene_cut_prefix", DEFAULT_SCENE_CUT_PREFIX) + single_video_only = getattr(config, "uniform_prompt", False) + allow_padding = getattr(config, "allow_padding", False) + min_latent_frames = getattr(config, "min_latent_frames", 0) + dataset_sample_warning_seconds = getattr(config, "dataset_sample_warning_seconds", 60.0) + dataset_sample_warning_interval_seconds = getattr( + config, "dataset_sample_warning_interval_seconds", 60.0 + ) + + if self.use_backward_simulation and not self.config.i2v: + dataset = MultiTextConcatDataset( + data_path=config.data_path, + num_blocks=num_blocks, + chunks_per_shot=chunks_per_shot, + scene_cut_prefix=scene_cut_prefix, + ) + collate_fn = eval_collate_fn + if dist.get_rank() == 0: + print(f"[backward_simulation] Using MultiTextConcatDataset: " + f"data_path={config.data_path}, num_blocks={num_blocks}, " + f"chunks_per_shot={chunks_per_shot}") + else: + dataset = MultiVideoConcatDataset( + data_dir=config.data_path, + video_size=(frame_raw_height, frame_raw_width), + total_frames=total_frames, + deterministic=False, + num_frame_per_block=num_frame_per_block, + temporal_compression_ratio=temporal_compression_ratio, + target_fps=self.fps, + allow_padding=allow_padding, + min_latent_frames=min_latent_frames, + single_video_only=single_video_only, + independent_first_frame=getattr(config, "independent_first_frame", False), + return_image=getattr(config, "i2v", False), + max_chunks_per_shot=chunks_per_shot, + scene_cut_prefix=scene_cut_prefix, + sample_warning_seconds=dataset_sample_warning_seconds, + sample_warning_interval_seconds=dataset_sample_warning_interval_seconds, + ) + collate_fn = multi_video_collate_fn + if dist.get_rank() == 0 and single_video_only: + print(f"[uniform_prompt] single_video_only enabled: each sample uses one video only") + random_seed = int(time.time()) % (2**31) * dist.get_rank() + sampler = torch.utils.data.distributed.DistributedSampler( + dataset, shuffle=True, drop_last=True, seed=random_seed) + dataloader = torch.utils.data.DataLoader( + dataset, batch_size=config.batch_size, sampler=sampler, + num_workers=2, prefetch_factor=1, pin_memory=False, + persistent_workers=False, collate_fn=collate_fn, + ) + + if dist.get_rank() == 0: + print("DATASET SIZE %d" % len(dataset)) + self.dataloader = cycle(dataloader) + + # Step 6: Initialize the validation dataloader for visualization (fixed prompts) + self.fixed_vis_batch = None + self.vis_interval = section_get(config, "evaluation", "interval", getattr(config, "vis_interval", -1)) + configured_vis_lengths = section_get(config, "evaluation", "num_frames", getattr(config, "vis_video_lengths", [])) + self.save_vis_latents_only = section_get( + config, + "evaluation", + "save_latents_only", + getattr(config, "return_latents", True), + aliases=("return_latents", "save_latent_only"), + ) + if isinstance(configured_vis_lengths, int): + configured_vis_lengths = [configured_vis_lengths] + if self.vis_interval > 0 and len(configured_vis_lengths) > 0: + # Determine validation data path + val_data_path = ( + getattr(config, "eval_data_path", None) + or getattr(config, "val_data_path", None) + or config.data_path + ) + + if self.config.i2v: + val_dataset = MultiVideoConcatDataset( + data_dir=val_data_path, + video_size=(frame_raw_height, frame_raw_width), + total_frames=total_frames, + deterministic=True, + num_frame_per_block=num_frame_per_block, + temporal_compression_ratio=temporal_compression_ratio, + target_fps=self.fps, + allow_padding=allow_padding, + min_latent_frames=min_latent_frames, + single_video_only=single_video_only, + independent_first_frame=getattr(config, "independent_first_frame", False), + return_image=True, + max_chunks_per_shot=chunks_per_shot, + scene_cut_prefix=scene_cut_prefix, + sample_warning_seconds=dataset_sample_warning_seconds, + sample_warning_interval_seconds=dataset_sample_warning_interval_seconds, + ) + val_collate_fn = multi_video_collate_fn + else: + val_dataset = MultiTextConcatDataset( + data_path=val_data_path, + num_blocks=num_blocks, + chunks_per_shot=chunks_per_shot, + scene_cut_prefix=scene_cut_prefix, + deterministic=True, + ) + val_collate_fn = eval_collate_fn + + if dist.get_rank() == 0: + print("VAL DATASET SIZE %d" % len(val_dataset)) + + sampler = torch.utils.data.distributed.DistributedSampler( + val_dataset, shuffle=False, drop_last=False) + val_dataloader = torch.utils.data.DataLoader( + val_dataset, + batch_size=section_get(config, "evaluation", "val_batch_size", getattr(config, "val_batch_size", 1)), + sampler=sampler, + num_workers=0, + collate_fn=val_collate_fn, + ) + + # Take the first batch as fixed visualization batch + try: + self.fixed_vis_batch = next(iter(val_dataloader)) + except StopIteration: + self.fixed_vis_batch = None + + # ---------------------------------------------------------------------------------------------------------- + # Visualization settings + # ---------------------------------------------------------------------------------------------------------- + # List of video lengths to visualize, e.g. [8, 16, 32] + self.vis_video_lengths = configured_vis_lengths + for _vl in self.vis_video_lengths: + assert _vl <= latent_frames_for_dataset, ( + f"vis_video_lengths entry {_vl} exceeds " + f"image_or_video_shape[1] ({latent_frames_for_dataset}), " + f"the dataset will not provide enough prompts for visualization." + ) + + if self.vis_interval > 0 and len(self.vis_video_lengths) > 0: + self._setup_visualizer() + + if not self.is_lora_enabled: + # ================================= Standard (non-LoRA) model logic ================================= + checkpoint_path = None + + if auto_resume and self.output_path: + # Auto resume: find latest checkpoint in logdir + latest_checkpoint = self.find_latest_checkpoint(self.output_path) + if latest_checkpoint: + checkpoint_path = latest_checkpoint + if self.is_main_process: + print(f"Auto resume: Found latest checkpoint at {checkpoint_path}") + else: + if self.is_main_process: + print("Auto resume: No checkpoint found in logdir, starting from scratch") + elif auto_resume: + if self.is_main_process: + print("Auto resume enabled but no logdir specified, starting from scratch") + else: + if self.is_main_process: + print("Auto resume disabled, starting from scratch") + + if checkpoint_path is None: + if getattr(config, "generator_ckpt", False): + # Explicit checkpoint path provided + checkpoint_path = config.generator_ckpt + if self.is_main_process: + print(f"Using explicit checkpoint: {checkpoint_path}") + + # Pre-load fake_score from a separate checkpoint if specified. + # This will be overwritten if checkpoint_path also contains "critic". + fake_score_ckpt = getattr(config, "fake_score_ckpt", None) + if fake_score_ckpt: + if self.is_main_process: + print(f"Loading fake_score from {fake_score_ckpt}") + fake_ckpt = torch.load(fake_score_ckpt, map_location="cpu") + if "critic" in fake_ckpt: + self.model.fake_score.load_state_dict(fake_ckpt["critic"], strict=True) + elif "fake_score" in fake_ckpt: + self.model.fake_score.load_state_dict(fake_ckpt["fake_score"], strict=True) + elif "model" in fake_ckpt: + self.model.fake_score.load_state_dict(fake_ckpt["model"], strict=True) + else: + if self.is_main_process: + print(f"No recognized key in {fake_score_ckpt}, treating as raw state_dict") + self.model.fake_score.load_state_dict(fake_ckpt, strict=True) + del fake_ckpt + gc.collect() + if self.is_main_process: + print(f"Successfully loaded fake_score from {fake_score_ckpt}") + + if checkpoint_path: + if self.is_main_process: + print(f"Loading checkpoint from {checkpoint_path}") + checkpoint = torch.load(checkpoint_path, map_location="cpu") + + # Load generator + if "generator" in checkpoint: + if self.is_main_process: + print(f"Loading pretrained generator from {checkpoint_path}") + self.model.generator.load_state_dict(checkpoint["generator"], strict=True) + del checkpoint["generator"] + elif "model" in checkpoint: + if self.is_main_process: + print(f"Loading pretrained generator from {checkpoint_path}") + self.model.generator.load_state_dict(checkpoint["model"], strict=True) + del checkpoint["model"] + else: + if self.is_main_process: + print(f"No 'generator'/'model' key found in {checkpoint_path}, treating as raw state_dict") + self.model.generator.load_state_dict(checkpoint, strict=True) + + # Load critic + if "critic" in checkpoint: + if self.is_main_process: + print(f"Loading pretrained critic from {checkpoint_path}") + self.model.fake_score.load_state_dict(checkpoint["critic"], strict=True) + del checkpoint["critic"] + else: + if self.is_main_process: + print("Warning: Critic checkpoint not found.") + + # Load EMA + if "generator_ema" in checkpoint and self.generator_ema is not None: + if self.is_main_process: + print(f"Loading pretrained EMA from {checkpoint_path}") + self.generator_ema.load_state_dict(checkpoint["generator_ema"]) + del checkpoint["generator_ema"] + else: + if self.is_main_process: + print("Warning: EMA checkpoint not found or EMA not initialized.") + + gc.collect() + + # For auto resume, always resume full training state + # Load optimizers + if "generator_optimizer" in checkpoint: + if self.is_main_process: + print("Resuming generator optimizer...") + gen_osd = FSDP.optim_state_dict_to_load( + self.model.generator, # FSDP root module + self.generator_optimizer, # newly created optimizer + checkpoint["generator_optimizer"] # optimizer state dict at save time + ) + del checkpoint["generator_optimizer"] + self.generator_optimizer.load_state_dict(gen_osd) + del gen_osd + else: + if self.is_main_process: + print("Warning: Generator optimizer checkpoint not found.") + + if "critic_optimizer" in checkpoint: + if self.is_main_process: + print("Resuming critic optimizer...") + crit_osd = FSDP.optim_state_dict_to_load( + self.model.fake_score, + self.critic_optimizer, + checkpoint["critic_optimizer"] + ) + del checkpoint["critic_optimizer"] + self.critic_optimizer.load_state_dict(crit_osd) + del crit_osd + else: + if self.is_main_process: + print("Warning: Critic optimizer checkpoint not found.") + + # Load training step + if "step" in checkpoint: + self.step = checkpoint["step"] + if self.is_main_process: + print(f"Resuming from step {self.step}") + else: + if self.is_main_process: + print("Warning: Step not found in checkpoint, starting from step 0.") + del checkpoint + gc.collect() + + # Load real_score from a separate checkpoint (independent of auto_resume) + real_score_ckpt = getattr(config, "real_score_ckpt", None) + if real_score_ckpt and not (self.real_score_quant and self.real_score_quant_materialize): + if self.is_main_process: + print(f"Loading real_score from {real_score_ckpt}") + real_ckpt = torch.load(real_score_ckpt, map_location="cpu") + if "generator" in real_ckpt: + self.model.real_score.load_state_dict(real_ckpt["generator"], strict=True) + elif "critic" in real_ckpt: + self.model.real_score.load_state_dict(real_ckpt["critic"], strict=True) + elif "model" in real_ckpt: + self.model.real_score.load_state_dict(real_ckpt["model"], strict=True) + else: + if self.is_main_process: + print(f"No recognized key in {real_score_ckpt}, treating as raw state_dict") + self.model.real_score.load_state_dict(real_ckpt, strict=True) + del real_ckpt + gc.collect() + if self.is_main_process: + print(f"Successfully loaded real_score from {real_score_ckpt}") + + ############################################################################################################## + + # Let's delete EMA params for early steps to save some computes at training and inference + # Note: This should be done after potential resume to avoid accidentally deleting resumed EMA + if self.step < config.ema_start_step: + self.generator_ema = None + + self.max_grad_norm_generator = getattr(config, "max_grad_norm_generator", 10.0) + self.max_grad_norm_critic = getattr(config, "max_grad_norm_critic", 10.0) + self.gradient_accumulation_steps = getattr(config, "gradient_accumulation_steps", 1) + self.previous_time = None + + if self.is_main_process: + print(f"Gradient accumulation steps: {self.gradient_accumulation_steps}") + if self.gradient_accumulation_steps > 1: + print(f"Effective batch size: {config.batch_size * self.gradient_accumulation_steps * self.world_size}") + + def _move_optimizer_to_device(self, optimizer, device): + """Move optimizer state to the specified device.""" + for state in optimizer.state.values(): + for k, v in state.items(): + if isinstance(v, torch.Tensor): + state[k] = v.to(device) + + def _materialize_quantized_model_before_fsdp( + self, + model, + model_label, + cache_transposed_weights=False, + ): + """Materialize NVFP4 weights before FSDP wraps quantized modules.""" + from utils.quant import _materialize_quantized_weights_for_inference + + current_rank = dist.get_rank() + target_device = torch.device("cuda", torch.cuda.current_device()) + if self.is_main_process: + print(f"[NVFP4] Materializing {model_label} sequentially before FSDP") + + for materialize_rank in range(self.world_size): + if current_rank == materialize_rank: + first_param = next(model.parameters(), None) + model_device = first_param.device if first_param is not None else target_device + if model_device != target_device: + model.to(target_device) + mat_modules, master_bytes, quant_bytes = _materialize_quantized_weights_for_inference( + model, + target_device=target_device, + cache_transposed_weights=cache_transposed_weights, + ) + if self.is_main_process: + print( + f"[NVFP4] {model_label} materialized: {len(mat_modules)} modules, " + f"master_weight={master_bytes / (1024**3):.3f} GiB freed, " + f"quantized_weight={quant_bytes / (1024**3):.3f} GiB" + ) + gc.collect() + dist.barrier() + + def find_latest_checkpoint(self, logdir): + """Find the latest checkpoint in the logdir.""" + if not os.path.exists(logdir): + return None + + checkpoint_dirs = [] + for item in os.listdir(logdir): + if item.startswith("checkpoint_model_") and os.path.isdir(os.path.join(logdir, item)): + try: + # Extract step number from directory name + step_str = item.replace("checkpoint_model_", "") + step = int(step_str) + checkpoint_path = os.path.join(logdir, item, "model.pt") + if os.path.exists(checkpoint_path): + checkpoint_dirs.append((step, checkpoint_path)) + except ValueError: + continue + + if not checkpoint_dirs: + return None + + # Sort by step number and return the latest one + checkpoint_dirs.sort(key=lambda x: x[0]) + latest_step, latest_path = checkpoint_dirs[-1] + return latest_path + + def get_all_checkpoints(self, logdir): + """Get all checkpoints in the logdir sorted by step number.""" + if not os.path.exists(logdir): + return [] + + checkpoint_dirs = [] + for item in os.listdir(logdir): + if item.startswith("checkpoint_model_") and os.path.isdir(os.path.join(logdir, item)): + try: + # Extract step number from directory name + step_str = item.replace("checkpoint_model_", "") + step = int(step_str) + checkpoint_dir_path = os.path.join(logdir, item) + checkpoint_file_path = os.path.join(checkpoint_dir_path, "model.pt") + if os.path.exists(checkpoint_file_path): + checkpoint_dirs.append((step, checkpoint_dir_path, item)) + except ValueError: + continue + + # Sort by step number (ascending order) + checkpoint_dirs.sort(key=lambda x: x[0]) + return checkpoint_dirs + + def cleanup_old_checkpoints(self, logdir, max_checkpoints): + """Remove old checkpoints if the number exceeds max_checkpoints. + + Only the main process performs the actual deletion to avoid race conditions + in distributed training. + """ + if max_checkpoints <= 0: + return + + # Only main process should perform cleanup to avoid race conditions + if not self.is_main_process: + return + + checkpoints = self.get_all_checkpoints(logdir) + if len(checkpoints) > max_checkpoints: + # Calculate how many to remove + num_to_remove = len(checkpoints) - max_checkpoints + checkpoints_to_remove = checkpoints[:num_to_remove] # Remove oldest ones + + print(f"Checkpoint cleanup: Found {len(checkpoints)} checkpoints, removing {num_to_remove} oldest ones (keeping {max_checkpoints})") + + import shutil + removed_count = 0 + for step, checkpoint_dir_path, dir_name in checkpoints_to_remove: + try: + print(f" Removing: {dir_name} (step {step})") + shutil.rmtree(checkpoint_dir_path) + removed_count += 1 + except Exception as e: + print(f" Warning: Failed to remove checkpoint {dir_name}: {e}") + + print(f"Checkpoint cleanup completed: removed {removed_count}/{num_to_remove} old checkpoints") + else: + if len(checkpoints) > 0: + print(f"Checkpoint cleanup: Found {len(checkpoints)} checkpoints (max: {max_checkpoints}, no cleanup needed)") + + def save(self): + print("Start gathering distributed model states...") + + if self.is_lora_enabled: + gen_lora_sd = self._gather_lora_state_dict( + self.model.generator.model) + crit_lora_sd = self._gather_lora_state_dict( + self.model.fake_score.model) + + state_dict = { + "generator_lora": gen_lora_sd, + "critic_lora": crit_lora_sd, + "step": self.step, + } + else: + with FSDP.state_dict_type( + self.model.generator, + StateDictType.FULL_STATE_DICT, + FullStateDictConfig(rank0_only=True, offload_to_cpu=True), + FullOptimStateDictConfig(rank0_only=True, offload_to_cpu=True), + ): + generator_state_dict = self.model.generator.state_dict() + generator_opim_state_dict = FSDP.optim_state_dict(self.model.generator, + self.generator_optimizer) + + if dist.is_initialized(): + dist.barrier() + + with FSDP.state_dict_type( + self.model.fake_score, + StateDictType.FULL_STATE_DICT, + FullStateDictConfig(rank0_only=True, offload_to_cpu=True), + FullOptimStateDictConfig(rank0_only=True, offload_to_cpu=True), + ): + critic_state_dict = self.model.fake_score.state_dict() + critic_opim_state_dict = FSDP.optim_state_dict(self.model.fake_score, + self.critic_optimizer) + + if self.config.ema_start_step < self.step and self.generator_ema is not None: + state_dict = { + "generator": generator_state_dict, + "critic": critic_state_dict, + "generator_ema": self.generator_ema.state_dict(), + "generator_optimizer": generator_opim_state_dict, + "critic_optimizer": critic_opim_state_dict, + "step": self.step, + } + else: + state_dict = { + "generator": generator_state_dict, + "critic": critic_state_dict, + "generator_optimizer": generator_opim_state_dict, + "critic_optimizer": critic_opim_state_dict, + "step": self.step, + } + + if self.is_main_process: + checkpoint_dir = os.path.join(self.output_path, f"checkpoint_model_{self.step:06d}") + os.makedirs(checkpoint_dir, exist_ok=True) + checkpoint_file = os.path.join(checkpoint_dir, "model.pt") + torch.save(state_dict, checkpoint_file) + print("Model saved to", checkpoint_file) + + # Cleanup old checkpoints if max_checkpoints is set + max_checkpoints = getattr(self.config, "max_checkpoints", 0) + if max_checkpoints > 0: + self.cleanup_old_checkpoints(self.output_path, max_checkpoints) + + # Keep all ranks in sync so non-rank0 workers don't kick off the next + # training iteration (and trigger NCCL watchdog timeouts) while rank0 + # is still writing the checkpoint to disk. + if dist.is_initialized(): + dist.barrier() + + torch.cuda.empty_cache() + import gc + gc.collect() + + def fwdbwd_one_step(self, batch, train_generator): + self.model.eval() # prevent any randomness (e.g. dropout) + + if self.step % 5 == 0: + torch.cuda.empty_cache() + + # Step 1: Get the next batch of text prompts + text_prompts = batch["prompts"] + + if getattr(self.config, "uniform_prompt", False): + text_prompts = [[sample[0]] * len(sample) for sample in text_prompts] + + batch_size = len(text_prompts) + image_or_video_shape = list(self.config.image_or_video_shape) + image_or_video_shape[0] = batch_size + + # Step 1.5: Prepare clean_latent and initial_latent for off-policy mode + clean_latent = None + initial_latent = None + if not self.use_backward_simulation: + if not getattr(self.config, "load_raw_video", False): + clean_latent = batch["ode_latent"][:, -1].to( + device=self.device, dtype=self.dtype) + else: + frames = batch["frames"].to( + device=self.device, dtype=self.dtype) + with torch.no_grad(): + clean_latent = self.model.vae.encode_to_latent(frames).to( + device=self.device, dtype=self.dtype) + if getattr(self.config, "i2v", False): + initial_latent = clean_latent[:, 0:1] + elif getattr(self.config, "i2v", False): + image = batch.get("image", None) + if image is None: + raise ValueError("DMD i2v backward-simulation requires batch['image'].") + image = image.to(device=self.device, dtype=self.dtype) + if image.ndim == 4: + image = image.unsqueeze(2) + elif image.ndim != 5: + raise ValueError( + f"Expected i2v image with shape [B,C,H,W] or [B,C,T,H,W], got {tuple(image.shape)}" + ) + with torch.no_grad(): + initial_latent = self.model.vae.encode_to_latent(image).to( + device=self.device, dtype=self.dtype) + + # Step 2: Extract the conditional infos + with torch.no_grad(): + # MultiVideoConcatDataset returns List[List[str]], flatten to List[str] + text_prompts_flat = [p for sublist in text_prompts for p in sublist] + conditional_dict = self.model.text_encoder( + text_prompts=text_prompts_flat) + + if not getattr(self, "unconditional_dict", None): + unconditional_dict = self.model.text_encoder( + text_prompts=[self.config.negative_prompt] * batch_size) + unconditional_dict = {k: v.detach() + for k, v in unconditional_dict.items()} + self.unconditional_dict = unconditional_dict + else: + unconditional_dict = self.unconditional_dict + + use_scene_cut_mask = ( + section_get(self.config, "inference", "multi_shot_sink", False) + or section_get( + self.config, + "inference", + "multi_shot_rope_offset", + 0.0, + ) != 0.0 + ) + if use_scene_cut_mask: + _prefix = getattr(self.config, "scene_cut_prefix", DEFAULT_SCENE_CUT_PREFIX) + conditional_dict["scene_cut_mask"] = [ + p.startswith(_prefix) for p in text_prompts[0] + ] + + # Step 3: Store gradients for the generator (if training the generator) + if train_generator: + generator_loss, generator_log_dict = self.model.generator_loss( + image_or_video_shape=image_or_video_shape, + conditional_dict=conditional_dict, + unconditional_dict=unconditional_dict, + clean_latent=clean_latent, + initial_latent=initial_latent + ) + + # Scale loss for gradient accumulation and backward + scaled_generator_loss = generator_loss / self.gradient_accumulation_steps + scaled_generator_loss.backward() + generator_log_dict.update({"generator_loss": generator_loss, + "generator_grad_norm": torch.tensor(0.0, device=self.device)}) + + return generator_log_dict + else: + generator_log_dict = {} + + critic_loss, critic_log_dict = self.model.critic_loss( + image_or_video_shape=image_or_video_shape, + conditional_dict=conditional_dict, + unconditional_dict=unconditional_dict, + clean_latent=clean_latent, + initial_latent=initial_latent + ) + + # Scale loss for gradient accumulation and backward + scaled_critic_loss = critic_loss / self.gradient_accumulation_steps + scaled_critic_loss.backward() + critic_log_dict.update({"critic_loss": critic_loss, + "critic_grad_norm": torch.tensor(0.0, device=self.device)}) + + return critic_log_dict + + def generate_video(self, pipeline, num_frames, prompts, image=None, latents_only=False): + batch_size = len(prompts) + if image is not None: + image = image.to(device=self.device, dtype=self.dtype) + if image.ndim == 4: + image = image.unsqueeze(2) + elif image.ndim != 5: + raise ValueError(f"Expected i2v image with shape [B,C,H,W] or [B,C,T,H,W], got {tuple(image.shape)}") + + # Encode the input image as the first latent + initial_latent = pipeline.vae.encode_to_latent(image).to(device=self.device, dtype=self.dtype) + if initial_latent.shape[0] != batch_size: + initial_latent = initial_latent.repeat(batch_size, 1, 1, 1, 1) + num_noise_frames = num_frames + if num_noise_frames <= initial_latent.shape[1]: + raise ValueError( + f"num_frames must exceed the i2v conditioning frames; " + f"got {num_frames} and {initial_latent.shape[1]}" + ) + sampled_noise = torch.randn( + [batch_size, num_noise_frames, self.config.image_or_video_shape[2], self.config.image_or_video_shape[3], self.config.image_or_video_shape[4]], + device=self.device, + dtype=self.dtype + ) + else: + initial_latent = None + sampled_noise = torch.randn( + [batch_size, num_frames, self.config.image_or_video_shape[2], self.config.image_or_video_shape[3], self.config.image_or_video_shape[4]], + device=self.device, + dtype=self.dtype + ) + with torch.no_grad(): + kwargs = dict(noise=sampled_noise, text_prompts=prompts) + if initial_latent is not None: + kwargs["initial_latent"] = initial_latent + kwargs["return_latents"] = latents_only + result = pipeline.inference(**kwargs) + if latents_only: + return result + video = result + current_video = video.permute(0, 1, 3, 4, 2).cpu().numpy() * 255.0 + if hasattr(pipeline, 'vae') and hasattr(pipeline.vae, 'model') and hasattr(pipeline.vae.model, 'clear_cache'): + pipeline.vae.model.clear_cache() + return current_video + + def train(self): + start_step = self.step + try: + while True: + # Check if we should train generator on this optimization step + TRAIN_GENERATOR = self.step % self.config.dfake_gen_update_ratio == 0 + + if TRAIN_GENERATOR: + self.generator_optimizer.zero_grad(set_to_none=True) + self.critic_optimizer.zero_grad(set_to_none=True) + + # Whole-cycle gradient accumulation loop + accumulated_generator_logs = [] + accumulated_critic_logs = [] + + for accumulation_step in range(self.gradient_accumulation_steps): + batch = next(self.dataloader) + + # Train generator (if needed) + if TRAIN_GENERATOR: + extra_gen = self.fwdbwd_one_step(batch, True) + accumulated_generator_logs.append(extra_gen) + + # Train critic + extra_crit = self.fwdbwd_one_step(batch, False) + accumulated_critic_logs.append(extra_crit) + + # Compute grad norm and update parameters + if TRAIN_GENERATOR: + generator_grad_norm = self.model.generator.clip_grad_norm_(self.max_grad_norm_generator) + generator_log_dict = merge_dict_list(accumulated_generator_logs) + generator_log_dict["generator_grad_norm"] = generator_grad_norm + + self.generator_optimizer.step() + if self.generator_ema is not None: + self.generator_ema.update(self.model.generator) + else: + generator_log_dict = {} + + critic_grad_norm = self.model.fake_score.clip_grad_norm_(self.max_grad_norm_critic) + critic_log_dict = merge_dict_list(accumulated_critic_logs) + critic_log_dict["critic_grad_norm"] = critic_grad_norm + + self.critic_optimizer.step() + + # Increment the step since we finished gradient update + self.step += 1 + + # Create EMA params (if not already created) + if (self.step >= self.config.ema_start_step) and \ + (self.generator_ema is None) and (self.config.ema_weight > 0): + if not self.is_lora_enabled: + self.generator_ema = EMA_FSDP(self.model.generator, decay=self.config.ema_weight) + if self.is_main_process: + print(f"EMA created at step {self.step} with weight {self.config.ema_weight}") + else: + if self.is_main_process: + print(f"EMA creation skipped at step {self.step} (disabled in LoRA mode)") + + # Save the model + if (not self.config.no_save) and (self.step - start_step) > 0 and self.step % self.config.log_iters == 0: + torch.cuda.empty_cache() + self.save() + torch.cuda.empty_cache() + + # Logging + if self.is_main_process: + wandb_loss_dict = {} + if TRAIN_GENERATOR and generator_log_dict: + wandb_loss_dict.update( + { + "generator_loss": generator_log_dict["generator_loss"].mean().item(), + "generator_grad_norm": generator_log_dict["generator_grad_norm"].mean().item(), + "dmdtrain_gradient_norm": generator_log_dict["dmdtrain_gradient_norm"].mean().item() + } + ) + + + wandb_loss_dict.update( + { + "critic_loss": critic_log_dict["critic_loss"].mean().item(), + "critic_grad_norm": critic_log_dict["critic_grad_norm"].mean().item() + } + ) + if not self.disable_wandb: + wandb.log(wandb_loss_dict, step=self.step) + + if self.step % self.config.gc_interval == 0: + if dist.get_rank() == 0: + logging.info("DistGarbageCollector: Running GC.") + gc.collect() + torch.cuda.empty_cache() + + if self.is_main_process: + current_time = time.time() + iteration_time = 0 if self.previous_time is None else current_time - self.previous_time + if not self.disable_wandb: + wandb.log({"per iteration time": iteration_time}, step=self.step) + self.previous_time = current_time + # Log training progress + if TRAIN_GENERATOR and generator_log_dict: + print(f"step {self.step}, per iteration time {iteration_time}, generator_loss {generator_log_dict['generator_loss'].mean().item()}, generator_grad_norm {generator_log_dict['generator_grad_norm'].mean().item()}, dmdtrain_gradient_norm {generator_log_dict['dmdtrain_gradient_norm'].mean().item()}, critic_loss {critic_log_dict['critic_loss'].mean().item()}, critic_grad_norm {critic_log_dict['critic_grad_norm'].mean().item()}") + else: + print(f"step {self.step}, per iteration time {iteration_time}, critic_loss {critic_log_dict['critic_loss'].mean().item()}, critic_grad_norm {critic_log_dict['critic_grad_norm'].mean().item()}") + + # ---------------------------------------- Visualization --------------------------------------------------- + + if self.vis_interval > 0 and (self.step % self.vis_interval == 0): + self._visualize() + + if self.step >= self.config.max_iters: + break + + except Exception as e: + print(f"[ERROR] [Rank {dist.get_rank()}] Training crashed at step {self.step} with exception: {e}") + print(f"[ERROR] [Rank {dist.get_rank()}] Exception traceback:", flush=True) + import traceback + traceback.print_exc() + + def _configure_lora_for_model(self, transformer, model_name): + """Configure LoRA for a WanDiffusionWrapper model""" + # Find all Linear modules in WanAttentionBlock modules + target_linear_modules = set() + + # Define the specific modules we want to apply LoRA to + all_causal = getattr(self.config, 'all_causal', False) + generator_is_causal = getattr(self.config, 'generator_is_causal', True) + if model_name == 'generator': + adapter_target_modules = ['CausalWanAttentionBlock'] if generator_is_causal else ['WanAttentionBlock'] + elif model_name == 'fake_score': + adapter_target_modules = ['CausalWanAttentionBlock'] if all_causal else ['WanAttentionBlock'] + else: + raise ValueError(f"Invalid model name: {model_name}") + + for name, module in transformer.named_modules(): + if module.__class__.__name__ in adapter_target_modules: + for full_submodule_name, submodule in module.named_modules(prefix=name): + if isinstance(submodule, torch.nn.Linear): + target_linear_modules.add(full_submodule_name) + + target_linear_modules = list(target_linear_modules) + + if self.is_main_process: + print(f"LoRA target modules for {model_name}: {len(target_linear_modules)} Linear layers") + if getattr(self.lora_config, 'verbose', False): + for module_name in sorted(target_linear_modules): + print(f" - {module_name}") + + # Create LoRA config + adapter_type = self.lora_config.get('type', 'lora') + if adapter_type == 'lora': + peft_config = peft.LoraConfig( + r=self.lora_config.get('rank', 16), + lora_alpha=self.lora_config.get('alpha', None) or self.lora_config.get('rank', 16), + lora_dropout=self.lora_config.get('dropout', 0.0), + target_modules=target_linear_modules, + ) + else: + raise NotImplementedError(f'Adapter type {adapter_type} is not implemented') + + # Apply LoRA to the transformer + lora_model = peft.get_peft_model(transformer, peft_config) + + if self.is_main_process: + print('peft_config', peft_config) + lora_model.print_trainable_parameters() + + return lora_model + + + def _gather_lora_state_dict(self, lora_model): + "On rank-0, gather FULL_STATE_DICT, then filter only LoRA weights" + with FSDP.state_dict_type( + lora_model, # lora_model contains nested FSDP submodules + StateDictType.FULL_STATE_DICT, + FullStateDictConfig(rank0_only=True, offload_to_cpu=True) + ): + full = lora_model.state_dict() + return get_peft_model_state_dict(lora_model, state_dict=full) + + # -------------------------------------------------------------------------------------------------------------- + # Visualization helpers + # -------------------------------------------------------------------------------------------------------------- + + def _setup_visualizer(self): + """Initialize the inference pipeline for visualization on CPU, to be moved to GPU only when needed.""" + + generator_is_causal = getattr(self.config, "generator_is_causal", True) + + if not generator_is_causal: + # Bidirectional generator: no pipeline object needed, + # _visualize_bidirectional handles the loop directly. + self.vis_pipeline = "bidirectional" + else: + from copy import deepcopy + vis_config = deepcopy(self.config) + if "guidance_scale" not in getattr(vis_config, "inference", {}): + vis_config.guidance_scale = 1.0 + if section_get(self.config, "inference", "sampling_steps", None) is None: + vis_config.sampling_steps = 50 + self.vis_pipeline = CausalDiffusionInferencePipeline( + args=vis_config, + device=self.device, + generator=self.model.generator, + text_encoder=self.model.text_encoder, + vae=self.model.vae) + + # Visualization output directory (default: /vis) + self.vis_output_dir = os.path.join(self.output_path, "vis") + os.makedirs(self.vis_output_dir, exist_ok=True) + if section_get(self.config, "evaluation", "use_ema", getattr(self.config, "vis_ema", False)): + raise NotImplementedError("Visualization with EMA is not implemented") + + @torch.no_grad() + def _generate_bidirectional(self, num_frames, prompts): + """Full-sequence bidirectional multi-step denoising for visualization.""" + from wan_5b.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler + + batch_size = len(prompts) + # Flatten prompts (List[List[str]] → List[str], take first per sample) + text_prompts_flat = [p[0] if isinstance(p, list) else p for p in prompts] + + conditional_dict = self.model.text_encoder(text_prompts=text_prompts_flat) + + noise = torch.randn( + [batch_size, num_frames, + self.config.image_or_video_shape[2], + self.config.image_or_video_shape[3], + self.config.image_or_video_shape[4]], + device=self.device, dtype=self.dtype) + + sampling_steps = section_get(self.config, "inference", "sampling_steps", getattr(self.config, "sampling_steps", 50)) + scheduler = self.model.generator.get_scheduler() + sample_scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=scheduler.num_train_timesteps, + shift=1, use_dynamic_shifting=False) + sample_scheduler.set_timesteps(sampling_steps, device=self.device, + shift=scheduler.shift) + + latents = noise + for t in sample_scheduler.timesteps: + timestep = t * torch.ones( + [batch_size, num_frames], device=self.device, dtype=torch.float32) + flow_pred, _ = self.model.generator( + noisy_image_or_video=latents, + conditional_dict=conditional_dict, + timestep=timestep, + ) + latents = sample_scheduler.step(flow_pred, t, latents, return_dict=False)[0] + + return latents + + def _visualize(self): + """Generate validation samples to monitor training progress.""" + if self.vis_interval <= 0 or not hasattr(self, "vis_pipeline"): + return False + + # FSDP forward includes communication, so every rank must enter + # visualization together; running rank 0 alone would hang. + + if not getattr(self, "fixed_vis_batch", None): + print("[Warning] No fixed validation batch available for visualization.") + return False + + step_vis_dir = os.path.join(self.vis_output_dir, f"step_{self.step:07d}") + os.makedirs(step_vis_dir, exist_ok=True) + batch = self.fixed_vis_batch + prompts = batch["prompts"] + + image = None + if self.config.i2v and ("image" in batch): + image = batch["image"] + + mode_info = "" + if self.is_lora_enabled: + mode_info = "_lora" + if self.is_main_process: + print(f"Generating latents in LoRA mode (step {self.step})") + + for vid_len in self.vis_video_lengths: + print(f"Generating validation samples of length {vid_len}") + if self.vis_pipeline == "bidirectional": + samples = self._generate_bidirectional(vid_len, prompts) + if not self.save_vis_latents_only: + samples = self.model.vae.decode_to_pixel(samples) + samples = (samples * 0.5 + 0.5).clamp(0, 1) + samples = samples.permute(0, 1, 3, 4, 2).cpu().numpy() * 255.0 + else: + samples = self.generate_video( + self.vis_pipeline, + vid_len, + prompts, + image=image, + latents_only=self.save_vis_latents_only, + ) + + for idx in range(samples.shape[0]): + if self.save_vis_latents_only: + sample_name = f"latents_step_{self.step:07d}_rank_{dist.get_rank()}_sample_{idx}_len_{vid_len}{mode_info}.pt" + out_path = os.path.join(step_vis_dir, sample_name) + torch.save(samples[idx].cpu(), out_path) + else: + sample_name = f"video_step_{self.step:07d}_rank_{dist.get_rank()}_sample_{idx}_len_{vid_len}{mode_info}.mp4" + out_path = os.path.join(step_vis_dir, sample_name) + write_video(out_path, torch.as_tensor(samples[idx]).to(torch.uint8), fps=24) + + del samples + torch.cuda.empty_cache() + + # Save prompts for reference + prompt_path = os.path.join( + step_vis_dir, + f"prompts_rank_{dist.get_rank()}.txt", + ) + with open(prompt_path, "w") as f: + for i, p in enumerate(prompts): + f.write(f"[sample {i}] {p}\n") + + # Release KV / cross-attention caches allocated during inference to prevent OOM + # when training resumes. These caches can consume ~20+ GB of GPU memory. + if hasattr(self.vis_pipeline, 'clear_cache'): + self.vis_pipeline.clear_cache() + + torch.cuda.empty_cache() + import gc + gc.collect() + + # Synchronize all ranks so that a crashed rank is detected immediately + # rather than causing a 10-minute NCCL timeout on the next training collective. + dist.barrier() + + return True diff --git a/trainer/sp_helper.py b/trainer/sp_helper.py new file mode 100644 index 0000000..b3f30ee --- /dev/null +++ b/trainer/sp_helper.py @@ -0,0 +1,18 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0 +# +# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Compatibility exports for LongLive SP training helpers.""" + +from wan_5b.distributed.sp_training import ( + DEFAULT_SP_VAE_HALO_LATENTS, + SequenceParallelHelper, +) + +__all__ = ["DEFAULT_SP_VAE_HALO_LATENTS", "SequenceParallelHelper"] diff --git a/utils/DATASET_PROMPT_LOGIC.md b/utils/DATASET_PROMPT_LOGIC.md new file mode 100644 index 0000000..37a0081 --- /dev/null +++ b/utils/DATASET_PROMPT_LOGIC.md @@ -0,0 +1,243 @@ +# Dataset Prompt 分配逻辑说明 + +本文档说明 `MultiTextConcatDataset`（纯文本）和 `MultiVideoConcatDataset`（视频训练）在不同配置下产生 `prompts` 的行为。 + +--- + +## 架构总览 + +| 使用场景 | Dataset 类 | 输出 | +|---|---|---| +| diffusion.py 训练 | `MultiVideoConcatDataset` | `frames` + `prompts` | +| diffusion.py 推理 | `MultiTextConcatDataset` | `prompts` | +| distillation.py 训练（backward_sim） | `MultiTextConcatDataset` | `prompts` | +| distillation.py 可视化 | `MultiTextConcatDataset` | `prompts` | +| inference.py | `MultiTextConcatDataset` | `prompts` | + +--- + +## 一、MultiTextConcatDataset（纯文本） + +### 关键参数 + +| 参数 | 含义 | +|---|---| +| `num_blocks` | 输出 prompt 列表的固定长度 | +| `chunks_per_shot` | 每个 shot 重复的 block 数（0=使用 even_durations 均分） | +| `scene_cut_prefix` | 切镜标记前缀，默认 `"The scene transitions. "` | + +简写约定：`0`, `1`, `2` 表示不同 caption；`p+X` 表示 `scene_cut_prefix + X` + +### 1. txt 模式（data_path 指向 .txt 文件） + +每行一个 caption，每个 sample 取 `idx` 行的 caption，重复 `num_blocks` 次。 +**不加** scene_cut_prefix（单 shot 语义）。 + +``` +data_path="prompts.txt", line[3]="A dog runs", num_blocks=12 +→ [A, A, A, A, A, A, A, A, A, A, A, A] +``` + +无论 `chunks_per_shot` 设为多少，txt 模式始终单 shot 重复。 + +### 2. 目录模式（data_path 指向目录） + +读取 `caption//*.json`（不需要 `video/` 目录）。 +每个 JSON 的 `caption` 字段作为一个 shot 的文本。 + +#### Shot duration 三级 fallback + +按优先级决定每个 shot 占多少个 block： + +1. **`shot_durations.txt`**（per-folder 文件）— 每行或逗号分隔的整数 +2. **`chunks_per_shot`**（全局 config）— 所有 shot 统一重复固定次数 +3. **`_even_durations`**（均分）— 将 `num_blocks` 均匀分给所有 caption + +#### 长度处理 + +输出始终恰好 `num_blocks` 个 prompt： +- **超过** → 截断尾部 +- **不足** → 用最后一个 caption 直接 padding（**不加** scene_cut_prefix） + +#### 示例 + +**3 caption, num_blocks=12, chunks_per_shot=4** + +``` +captions: [0, 1, 2] +shot_durations: [4, 4, 4] +→ [0, 0, 0, 0, p+1, 1, 1, 1, p+2, 2, 2, 2] +``` + +**3 caption, num_blocks=12, chunks_per_shot=0（even_durations）** + +``` +captions: [0, 1, 2] +even_durations: base=4, extra=0 → [4, 4, 4] +→ [0, 0, 0, 0, p+1, 1, 1, 1, p+2, 2, 2, 2] +``` + +**2 caption, num_blocks=12, chunks_per_shot=4** + +``` +captions: [0, 1] +shot_durations: [4, 4], sum=8 < 12 → 最后一个 shot 扩展到 8 +→ [0, 0, 0, 0, p+1, 1, 1, 1, 1, 1, 1, 1] + ↑ padding 不加 prefix +``` + +**5 caption, num_blocks=12, chunks_per_shot=4** + +``` +captions: [0, 1, 2, 3, 4] +shot_durations: [4, 4, 4, 4, 4], 但 num_blocks=12 → clamped 到 [4, 4, 4] +→ [0, 0, 0, 0, p+1, 1, 1, 1, p+2, 2, 2, 2] + caption 3 和 4 被截断 +``` + +**2 caption, num_blocks=12, chunks_per_shot=0（even_durations）** + +``` +captions: [0, 1] +even_durations: base=6, extra=0 → [6, 6] +→ [0, 0, 0, 0, 0, 0, p+1, 1, 1, 1, 1, 1] +``` + +**1 caption, num_blocks=12** + +``` +captions: [0] +→ [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] + 单 shot，不加 prefix +``` + +**shot_durations.txt 覆盖** + +``` +captions: [0, 1, 2], shot_durations.txt 内容: "2, 6, 4", num_blocks=12 +→ [0, 0, p+1, 1, 1, 1, 1, 1, p+2, 2, 2, 2] +``` + +--- + +## 二、MultiVideoConcatDataset（视频训练） + +### 关键参数 + +| 参数 | 含义 | +|---|---| +| `total_segments` | 总 segment 数量（= 1 + num_subsequent_segments） | +| `single_video_only` | config 中的 `uniform_prompt`，True 时只从一个视频采样 | +| `max_chunks_per_shot` | 单镜头最大连续 chunk 数，超过则跳 1 秒做虚拟切镜（0=不限制） | +| `scene_cut_prefix` | 切镜标记前缀 | +| `allow_padding` | 视频不够时是否允许 padding（否则跳过该 folder） | + +Prompt 由实际视频采样决定，逐 segment 从视频文件加载对应的 per-video caption。 + +### 1. 多视频自然拼接（`max_chunks_per_shot=0`，默认） + +按视频文件顺序采样，切换视频文件时加 `scene_cut_prefix`： + +``` +video A 够采 3 chunks, video B 够采 4 chunks, total_segments=7 +→ [A, A, A, p+B, B, B, B] +``` + +### 2. `single_video_only=True` + +强制只从一个视频文件采样。视频不够长则整个 folder 被跳过： + +``` +video A 够采 7 chunks, total_segments=7 +→ [A, A, A, A, A, A, A] + +video A 只够采 5 chunks, total_segments=7 +→ (失败，跳到下一个 folder) +``` + +### 3. `max_chunks_per_shot=3`（限制单镜头时长） + +从同一视频连续采超过 3 chunks 后，跳 1 秒做虚拟切镜，加 `scene_cut_prefix`： + +``` +video A 很长, video B, total_segments=7, max_chunks_per_shot=3 +→ [A, A, A, p+A, A, A, p+B] + ↑跳1秒虚拟切镜 ↑换视频 +``` + +如果跳 1 秒后 A 不够了，直接跳到 B： + +``` +video A 够 4 chunks（跳1秒后不够）, video B, total_segments=7, max_chunks_per_shot=3 +→ [A, A, A, p+B, B, B, B] + ↑A跳1秒后不够，换B +``` + +### 4. `single_video_only=True` + `max_chunks_per_shot=3` + +单视频内也可以做虚拟切镜： + +``` +video A 很长, total_segments=7, single_video_only=True, max_chunks_per_shot=3 +→ [A, A, A, p+A, A, A, p+A] +``` + +### 5. 训练 padding（`allow_padding=True`） + +视频不够时用最后一个 caption 直接 padding（不加 prefix）： + +``` +video A 够采 3 chunks, video B 够采 2 chunks, total_segments=7, allow_padding=True +→ [A, A, A, p+B, B, B, B] + ↑后 2 个用最后一个 caption padding +``` + +--- + +## 三、Multi-Shot Sink + +通过 config 的 `multi_shot_sink: true` 开启。 +开启后，当检测到某个 block 处于新场景的起始位置时，会在该 block 去噪完成并更新 cache 后， +将 KV cache 的 attention sink 从旧场景的第一帧迁移到新场景的第一帧。 +同时会自动把全局 sink 长度设为 `sink_size`，不需要单独配置 global sink。 + +**配置方式**（yaml config 中添加）： + +```yaml +sink_size: 8 +multi_shot_sink: true # 默认 false，不迁移 sink +``` + +此选项在以下所有场景均生效： + +| 场景 | Pipeline | 检测方式 | +|---|---|---| +| **Diffusion trainer evaluation** | `CausalDiffusionInferencePipeline` | 检查 prompt 是否以 `scene_cut_prefix` 开头 | +| **Distillation trainer evaluation** | `CausalDiffusionInferencePipeline` | 同上 | +| **离线推理 `inference.py`** | `CausalDiffusionInferencePipeline` | 同上 | +| **Distillation backward simulation** | `SelfForcingTrainingPipeline` | trainer 预计算 `scene_cut_mask` 传入 `conditional_dict` | +| **Streaming long tuning** | `SelfForcingTrainingPipeline` | 同上（mask 随 chunk 自动 slice） | + +### 实现机制 + +**Inference pipeline**（`CausalDiffusionInferencePipeline`）： +直接检查每个 chunk 的 raw prompt 是否以 `scene_cut_prefix` 开头。 + +**Training pipeline**（`SelfForcingTrainingPipeline`）： +由于 prompts 在进入 pipeline 前已编码为 embedding，无法从 embedding 反推文本。 +因此 trainer 在编码前从原始 prompts 计算布尔列表 `scene_cut_mask`， +放入 `conditional_dict["scene_cut_mask"]` 一路透传到 pipeline。 +对于 streaming training，`scene_cut_mask` 在 `_slice_cond_dict_for_chunk` 中随 prompt 一起 slice。 + +``` +block: [0] [1] [2] [p+3] [4] [5] [p+6] [7] +mask: F F F T F F T F +sink: 0 0 0 →3 3 3 →6 6 + ↑更新sink ↑更新sink +``` + +当 `multi_shot_sink: false`（默认）时，sink 始终锚定在视频的第一帧，不做任何迁移。 + +`scene_cut_prefix` 由 `DEFAULT_SCENE_CUT_PREFIX` 常量定义（`dataset.py`）， +inference pipeline 引用同一常量，确保训练和推理的切镜检测一致。 +可通过 config 中的 `scene_cut_prefix` 字段统一覆盖。 diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000..f756774 --- /dev/null +++ b/utils/__init__.py @@ -0,0 +1,5 @@ +# Marker file: turn `utils/` from a namespace package into a regular package. +# torch 2.12's torchrun + multiprocessing has trouble resolving namespace +# packages from cwd in subprocesses; making this an explicit regular package +# makes `from utils.position_embedding_utils import ...` (used inside +# wan_5b/modules/model.py) reliable across torch versions. diff --git a/utils/adaln_triton.py b/utils/adaln_triton.py new file mode 100644 index 0000000..ab0a40d --- /dev/null +++ b/utils/adaln_triton.py @@ -0,0 +1,126 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0 +# +# SPDX-License-Identifier: Apache-2.0 +"""Triton fused adaLN-modulation kernel. + +iter-43: replaces the + (norm(x).unflatten(1, (F, frame_seqlen)) * (1 + e_scale) + e_shift).flatten(1, 2) +chain (LayerNorm + 2 broadcast elementwise ops per call, x2 per transformer block +for norm1/norm2) with a single Triton kernel. Each token does one fp32 pass: +mean/var reduce, normalize, multiply by (1+scale), add shift, cast back. + +LayerNorm matches `WanLayerNorm` (nn.LayerNorm with elementwise_affine=False, +eps=1e-6, output cast back to input dtype). +""" +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +# iter-44: autotune over num_warps + num_stages. Fixed BLOCK_C +# (next_power_of_2(C)) — varying it would change the reduce semantics. Provide +# enough configs to cover the (small_B, large_L) regime of inference chunks. +_ADALN_CONFIGS = [ + triton.Config({}, num_warps=nw, num_stages=ns) + for nw in (4, 8, 16) + for ns in (1, 2, 3) +] + + +@triton.autotune(configs=_ADALN_CONFIGS, key=["C", "FRAME_SEQLEN"]) +@triton.jit +def _adaln_modulate_kernel( + x_ptr, # [B, L, C] + scale_ptr, # [B, F, 1, C] (or any layout, indexed via strides) + shift_ptr, # [B, F, 1, C] + out_ptr, # [B, L, C] + C, + FRAME_SEQLEN, + x_stride_b, x_stride_l, + scale_stride_b, scale_stride_f, + shift_stride_b, shift_stride_f, + eps, + ADD_ONE_TO_SCALE: tl.constexpr, + BLOCK_C: tl.constexpr, +): + pid_b = tl.program_id(0) + pid_l = tl.program_id(1) + pid_f = pid_l // FRAME_SEQLEN + + offs_c = tl.arange(0, BLOCK_C) + mask = offs_c < C + + x_off = pid_b * x_stride_b + pid_l * x_stride_l + offs_c + x = tl.load(x_ptr + x_off, mask=mask, other=0.0).to(tl.float32) + + inv_C = 1.0 / C + mean = tl.sum(x, axis=0) * inv_C + x_centered = tl.where(mask, x - mean, 0.0) + var = tl.sum(x_centered * x_centered, axis=0) * inv_C + rstd = 1.0 / tl.sqrt(var + eps) + # Match WanLayerNorm: nn.LayerNorm casts back to input dtype via .type_as(x) + # before downstream ops. Round-trip through bf16 to keep numerics identical + # to the eager path so latent diff stays in run-to-run noise floor. + x_norm = (x_centered * rstd).to(tl.bfloat16).to(tl.float32) + + s_off = pid_b * scale_stride_b + pid_f * scale_stride_f + offs_c + h_off = pid_b * shift_stride_b + pid_f * shift_stride_f + offs_c + scale = tl.load(scale_ptr + s_off, mask=mask, other=0.0).to(tl.float32) + shift = tl.load(shift_ptr + h_off, mask=mask, other=0.0).to(tl.float32) + + if ADD_ONE_TO_SCALE: + # Match eager: `1 + e_scale` happens in bf16 (autocast off at this site) + scale = (scale + 1.0) + scale = scale.to(tl.bfloat16).to(tl.float32) + + # Match eager: bf16 * bf16, bf16 + bf16 + prod = (x_norm * scale).to(tl.bfloat16).to(tl.float32) + y = prod + shift + + tl.store(out_ptr + x_off, y, mask=mask) + + +def adaln_modulate_triton( + x: torch.Tensor, # [B, L, C] contiguous + e_scale: torch.Tensor, # [B, F, 1, C] + e_shift: torch.Tensor, # [B, F, 1, C] + frame_seqlen: int, + eps: float = 1e-6, + add_one_to_scale: bool = True, +) -> torch.Tensor: + """Fused (LayerNorm + (1+e_scale)*x + e_shift) over [B, F*frame_seqlen, C]. + + Replaces the WanAttentionBlock norm1/norm2 + modulate pattern. Output dtype + follows x.dtype; internal arithmetic is fp32. eps matches `WanLayerNorm`. + """ + assert x.dim() == 3, f"x must be [B, L, C], got {x.shape}" + B, L, C = x.shape + assert L % frame_seqlen == 0, (L, frame_seqlen) + F = L // frame_seqlen + assert e_scale.dim() == 4 and e_scale.shape[0] == B and e_scale.shape[1] == F \ + and e_scale.shape[2] == 1 and e_scale.shape[3] == C, \ + f"e_scale {tuple(e_scale.shape)} vs expected {(B, F, 1, C)}" + assert e_shift.shape == e_scale.shape + + out = torch.empty_like(x) + BLOCK_C = triton.next_power_of_2(C) + + grid = (B, L) + # num_warps / num_stages picked by @triton.autotune (iter-44). + _adaln_modulate_kernel[grid]( + x, e_scale, e_shift, out, + C, frame_seqlen, + x.stride(0), x.stride(1), + e_scale.stride(0), e_scale.stride(1), + e_shift.stride(0), e_shift.stride(1), + eps, + ADD_ONE_TO_SCALE=add_one_to_scale, + BLOCK_C=BLOCK_C, + ) + return out diff --git a/utils/config.py b/utils/config.py new file mode 100644 index 0000000..068be5c --- /dev/null +++ b/utils/config.py @@ -0,0 +1,174 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0 +# +# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied. +# +# SPDX-License-Identifier: Apache-2.0 + +from omegaconf import OmegaConf + + +DEFAULT_NEGATIVE_PROMPT = ( + "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，" + "整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，" + "画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，" + "手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" +) + + +wan_default_config = { + "Wan2.2-TI2V-5B": { + "resolution": [1280, 704], + "temporal_compression_ratio": 4, + "spatial_compression_ratio": 16, + "num_heads": 24, + "head_dim": 128, + "num_transformer_blocks": 30, + "fps": 24, + } +} + + +SECTION_KEYS = ( + "infra", + "algorithm", + "training", + "data", + "evaluation", + "inference", + "logging", + "checkpoints", +) + + +def _set_once(config, key, value, source): + if value is None: + return + if key in config and config[key] != value: + raise ValueError( + f"{key} is defined more than once with different values: " + f"{config[key]} vs {value} from {source}." + ) + config[key] = value + + +def section_get(config, section_key, key, default=None, aliases=()): + """Read a grouped config value, falling back to legacy flat names.""" + section = config.get(section_key, None) + candidate_keys = (key, *aliases) + if section is not None: + for candidate in candidate_keys: + if candidate in section: + return section[candidate] + for candidate in candidate_keys: + if candidate in config: + return config[candidate] + return default + + +def normalize_config(config): + """Expand grouped release configs into the flat runtime schema. + + The training and inference code historically reads fields such as + ``config.batch_size`` and ``config.model_kwargs`` directly. Release + configs can group those fields for readability, then call this function at + the entry point to preserve the existing runtime contract. + """ + for section_key in SECTION_KEYS: + section = config.get(section_key, None) + if section is None: + continue + for key, value in section.items(): + config[key] = value + + evaluation = config.get("evaluation", None) + if evaluation is not None: + if "interval" in evaluation: + _set_once(config, "generate_interval", evaluation.interval, "evaluation.interval") + _set_once(config, "vis_interval", evaluation.interval, "evaluation.interval") + if "num_frames" in evaluation: + num_frames = evaluation.num_frames + if isinstance(num_frames, (list, tuple)): + vis_lengths = list(num_frames) + inference_num_frames = vis_lengths[0] if vis_lengths else 0 + else: + inference_num_frames = int(num_frames) + vis_lengths = [inference_num_frames] + _set_once(config, "inference_num_frames", inference_num_frames, "evaluation.num_frames") + _set_once(config, "vis_video_lengths", vis_lengths, "evaluation.num_frames") + if "use_ema" in evaluation: + _set_once(config, "vis_ema", evaluation.use_ema, "evaluation.use_ema") + + model_section = config.get("model", None) + base_model_kwargs = config.get("model_kwargs", None) + model_kwargs = OmegaConf.create({}) + if base_model_kwargs is not None: + model_kwargs = OmegaConf.merge(model_kwargs, base_model_kwargs) + + if model_section is not None: + section_kwargs = model_section.get("kwargs", None) + if section_kwargs is not None: + model_kwargs = OmegaConf.merge(model_kwargs, section_kwargs) + + model_name = model_section.get("name", None) + if model_name is not None: + model_kwargs.model_name = model_name + config.model_name = model_name + + _set_once( + config, + "num_frame_per_block", + model_section.get("num_frame_per_block", None), + "model.num_frame_per_block", + ) + + if "model_name" in config and "model_name" not in model_kwargs: + model_kwargs.model_name = config.model_name + if "timestep_shift" in config and "timestep_shift" not in model_kwargs: + model_kwargs.timestep_shift = config.timestep_shift + if "timestep_shift" in model_kwargs: + _set_once(config, "timestep_shift", model_kwargs.timestep_shift, "model_kwargs.timestep_shift") + + model_num_frame_per_block = model_kwargs.get("num_frame_per_block", None) + if model_num_frame_per_block is not None: + _set_once(config, "num_frame_per_block", model_num_frame_per_block, "model_kwargs.num_frame_per_block") + + if len(model_kwargs) > 0: + config.model_kwargs = model_kwargs + + if "wandb_host" not in config: + config.wandb_host = "https://api.wandb.ai" + + if "negative_prompt" not in config: + config.negative_prompt = DEFAULT_NEGATIVE_PROMPT + + if config.get("trainer", None) == "score_distillation": + dmd_defaults = { + "i2v": False, + "teacher_forcing": False, + "backward_simulation": True, + "independent_first_frame": False, + "num_train_timestep": 1000, + "denoising_loss_type": "flow", + "real_guidance_scale": 3.0, + "fake_guidance_scale": 0.0, + } + for key, value in dmd_defaults.items(): + if key not in config: + config[key] = value + if "causal" not in config: + config.causal = bool(config.get("all_causal", True)) + + # Causal DMD uses the same Wan backbone for generator/teacher/critic unless + # a role-specific override is explicitly provided. + if getattr(config, "all_causal", False) and "model_kwargs" in config: + for role_key in ("real_model_kwargs", "fake_model_kwargs"): + if config.get(role_key, None) is None: + config[role_key] = OmegaConf.create( + OmegaConf.to_container(config.model_kwargs, resolve=True) + ) + + return config diff --git a/utils/dataset.py b/utils/dataset.py new file mode 100644 index 0000000..ebddca0 --- /dev/null +++ b/utils/dataset.py @@ -0,0 +1,1073 @@ +# Adopted from https://github.com/guandeh17/Self-Forcing +# SPDX-License-Identifier: Apache-2.0 +from torch.utils.data import Dataset +import numpy as np +import torch +import random +import json +from pathlib import Path +from PIL import Image +import os +import subprocess +import time +import warnings +import torchvision.transforms as transforms +import torchvision.transforms.functional as F + +try: + import decord +except ModuleNotFoundError: + decord = None + +DEFAULT_SCENE_CUT_PREFIX = "The scene transitions. " + + +class TextDataset(Dataset): + def __init__(self, prompt_path, extended_prompt_path=None): + with open(prompt_path, encoding="utf-8") as f: + self.prompt_list = [line.rstrip() for line in f] + + if extended_prompt_path is not None: + with open(extended_prompt_path, encoding="utf-8") as f: + self.extended_prompt_list = [line.rstrip() for line in f] + assert len(self.extended_prompt_list) == len(self.prompt_list) + else: + self.extended_prompt_list = None + + def __len__(self): + return len(self.prompt_list) + + def __getitem__(self, idx): + batch = { + "prompts": self.prompt_list[idx], + "idx": idx, + } + if self.extended_prompt_list is not None: + batch["extended_prompts"] = self.extended_prompt_list[idx] + return batch + + +class MultiTextDataset(Dataset): + """Dataset for multi-segment prompts stored in a JSONL file. + + Each line is a JSON object, e.g. + {"prompts": ["a cat", "a dog", "a bird"]} + + Args + ---- + prompt_path : str + Path to the JSONL file + field : str + Name of the list-of-strings field, default "prompts" + cache_dir : str | None + ``cache_dir`` passed to HF Datasets (optional) + """ + + def __init__(self, prompt_path: str, field: str = "prompts", cache_dir: str | None = None): + try: + import datasets + except ModuleNotFoundError as exc: + raise ModuleNotFoundError( + "The 'datasets' package is required for MultiTextDataset. " + "Use MultiTextConcatDataset for plain txt/json-caption directories " + "or install datasets." + ) from exc + self.ds = datasets.load_dataset( + "json", + data_files=prompt_path, + split="train", + cache_dir=cache_dir, + streaming=False, + ) + + assert len(self.ds) > 0, "JSONL is empty" + assert field in self.ds.column_names, f"Missing field '{field}'" + + seg_len = len(self.ds[0][field]) + for i, ex in enumerate(self.ds): + val = ex[field] + assert isinstance(val, list), f"Line {i} field '{field}' is not a list" + assert len(val) == seg_len, f"Line {i} list length mismatch" + + self.field = field + + def __len__(self): + return len(self.ds) + + def __getitem__(self, idx: int): + return { + "idx": idx, + "prompts_list": self.ds[idx][self.field], # List[str] + } + + +class MultiTextConcatDataset(Dataset): + """Text-only dataset for multi-shot training and inference. + + Supports two input modes: + + **txt file** — each line is one caption. Each sample uses the caption at + index ``idx``, repeated ``num_blocks`` times (single-shot, no scene cut + prefix). + + **directory** — reads ``caption//*.json`` files (no video dir + needed). Shot durations are resolved with a three-level fallback: + + 1. ``shot_durations.txt`` in the caption subfolder (per-sample override) + 2. ``chunks_per_shot`` from config (global fixed repeat) + 3. Even distribution across all available captions + + Scene cut prefix is prepended at shot boundaries (first block of each + shot except shot 0). Output is always exactly ``num_blocks`` prompts: + truncated if too many, padded with the last caption if too few. + """ + + def __init__( + self, + data_path: str, + num_blocks: int, + chunks_per_shot: int = 0, + scene_cut_prefix: str = DEFAULT_SCENE_CUT_PREFIX, + caption_field: str = "caption", + deterministic: bool = False, + ): + self.num_blocks = num_blocks + self.chunks_per_shot = chunks_per_shot + self.scene_cut_prefix = scene_cut_prefix + self.caption_field = caption_field + self.deterministic = deterministic + + path = Path(data_path) + if data_path.endswith(".txt") or path.is_file(): + self._mode = "txt" + with open(data_path, encoding="utf-8") as f: + self._prompts = [line.rstrip() for line in f if line.strip()] + assert len(self._prompts) > 0, f"No prompts found in {data_path}" + else: + self._mode = "dir" + self._caption_dir = path / "caption" if (path / "caption").is_dir() else path + self._folders = sorted([d for d in self._caption_dir.iterdir() if d.is_dir()]) + assert len(self._folders) > 0, ( + f"No caption subfolders found in {self._caption_dir}" + ) + + def __len__(self): + if self._mode == "txt": + return len(self._prompts) + return len(self._folders) + + def __getitem__(self, idx): + if self._mode == "txt": + return self._get_txt_item(idx) + return self._get_dir_item(idx) + + # ------------------------------------------------------------------ + # txt mode + # ------------------------------------------------------------------ + + def _get_txt_item(self, idx): + caption = self._prompts[idx % len(self._prompts)] + return { + "prompts": [caption] * self.num_blocks, + "idx": idx, + } + + # ------------------------------------------------------------------ + # directory mode + # ------------------------------------------------------------------ + + def _get_dir_item(self, idx): + folder = self._folders[idx % len(self._folders)] + raw_captions = self._load_captions_from_folder(folder) + if not raw_captions: + raw_captions = [""] + + shot_durations = self._resolve_shot_durations(folder, len(raw_captions)) + prompts = self._apply_shot_durations(raw_captions, shot_durations) + + # Ensure exactly num_blocks prompts + if len(prompts) > self.num_blocks: + prompts = prompts[: self.num_blocks] + elif len(prompts) < self.num_blocks: + last = prompts[-1] if prompts else "" + prompts.extend([last] * (self.num_blocks - len(prompts))) + + return { + "prompts": prompts, + "idx": idx, + } + + def _load_captions_from_folder(self, folder: Path): + json_files = sorted( + [f for f in folder.glob("*.json") if f.name != "global.json"], + key=lambda p: (p.stem.isdigit(), int(p.stem) if p.stem.isdigit() else 0, p.stem), + ) + captions = [] + for jf in json_files: + try: + with open(jf, "r", encoding="utf-8") as f: + data = json.load(f) + captions.append(data.get(self.caption_field, "")) + except Exception: + captions.append("") + return captions + + # ------------------------------------------------------------------ + # shot duration helpers + # ------------------------------------------------------------------ + + @staticmethod + def _load_shot_durations(folder: Path): + txt_path = folder / "shot_durations.txt" + if not txt_path.exists(): + return None + try: + with open(txt_path, "r") as f: + content = f.read().strip() + parts = content.replace(",", " ").split() + durations = [int(x) for x in parts if x.strip()] + return durations if durations else None + except Exception: + return None + + def _resolve_shot_durations(self, folder: Path, num_captions: int): + durations = self._load_shot_durations(folder) + if durations is not None: + return durations[:num_captions] + if self.chunks_per_shot > 0: + return [self.chunks_per_shot] * num_captions + return self._even_durations(num_captions) + + def _even_durations(self, num_shots: int): + total = self.num_blocks + base, extra = divmod(total, num_shots) + return [base + (1 if i < extra else 0) for i in range(num_shots)] + + def _apply_shot_durations(self, raw_captions, shot_durations): + target = self.num_blocks + clamped: list[int] = [] + remaining = target + for d in shot_durations: + if remaining <= 0: + break + take = min(d, remaining) + clamped.append(take) + remaining -= take + if remaining > 0 and clamped: + clamped[-1] += remaining + + prompts: list[str] = [] + for shot_idx, (caption, duration) in enumerate(zip(raw_captions, clamped)): + for block_in_shot in range(duration): + if shot_idx > 0 and block_in_shot == 0 and self.scene_cut_prefix: + prompts.append(self.scene_cut_prefix + caption) + else: + prompts.append(caption) + return prompts + + +class MultiVideoConcatDataset(Dataset): + """Dataset that concatenates multiple videos from a folder into a fixed-length video. + + Each item consists of multiple video segments concatenated together: + - First segment: first_chunk_frames frames (chunk) + - Subsequent segments: subsequent_chunk_frames frames each (chunk) + - Total: total_frames frames (first_chunk_frames + subsequent_chunk_frames*num_subsequent_segments = total_frames) + + Videos are sampled preserving original duration, and if a video doesn't have enough frames, + it moves to the next video. If a video has enough frames, it can be sampled repeatedly. + """ + def __init__( + self, + data_dir, + video_size, + total_frames, + target_fps=16, + video_extensions=('.mp4', '.avi', '.mov', '.mkv', '.webm'), + caption_field='caption', + filter_invalid_folders=False, + deterministic: bool = False, + num_frame_per_block=8, + temporal_compression_ratio=4, + allow_padding: bool = False, + min_latent_frames: int = 0, + single_video_only: bool = False, + independent_first_frame: bool = False, + return_image: bool = False, + max_chunks_per_shot: int = 0, + scene_cut_prefix: str = DEFAULT_SCENE_CUT_PREFIX, + sample_warning_seconds: float = 60.0, + sample_warning_interval_seconds: float = 60.0, + ): + self.root_dir = Path(data_dir) + self.data_dir = self.root_dir / "video" + self.caption_dir = self.root_dir / "caption" + self.video_size = video_size + self.total_frames = total_frames + + total_latent_frames = 1 + (total_frames - 1) // temporal_compression_ratio + separate_first_latent = ( + independent_first_frame + and total_latent_frames % num_frame_per_block != 0 + ) + if separate_first_latent: + assert (total_latent_frames - 1) % num_frame_per_block == 0, ( + f"total latent frames ({total_latent_frames}) must be divisible by " + f"num_frame_per_block ({num_frame_per_block}) or equal to " + f"1 + N * num_frame_per_block when independent_first_frame=True" + ) + first_chunk_latent_frames = ( + num_frame_per_block + 1 if separate_first_latent else num_frame_per_block + ) + first_chunk_frames = 1 + (first_chunk_latent_frames - 1) * temporal_compression_ratio + subsequent_chunk_frames = num_frame_per_block * temporal_compression_ratio + + self.first_chunk_frames = first_chunk_frames + self.subsequent_chunk_frames = subsequent_chunk_frames + self.target_fps = target_fps + self.caption_field = caption_field + self.video_extensions = video_extensions + self.filter_invalid_folders = filter_invalid_folders + self.deterministic = deterministic + self.allow_padding = allow_padding + self.num_frame_per_block = num_frame_per_block + self.independent_first_frame = independent_first_frame + self.first_chunk_latent_frames = first_chunk_latent_frames + self.return_image = return_image + if min_latent_frames > 0: + assert min_latent_frames % num_frame_per_block == 0, ( + f"min_latent_frames ({min_latent_frames}) must be a multiple of " + f"num_frame_per_block ({num_frame_per_block})" + ) + self.min_latent_frames = min_latent_frames + self.single_video_only = single_video_only + self.max_chunks_per_shot = max_chunks_per_shot + self.scene_cut_prefix = scene_cut_prefix + self.sample_warning_seconds = float(sample_warning_seconds or 0.0) + self.sample_warning_interval_seconds = float(sample_warning_interval_seconds or 0.0) + + remaining_frames = total_frames - first_chunk_frames + self.num_subsequent_segments = remaining_frames // subsequent_chunk_frames + self.total_segments = 1 + self.num_subsequent_segments + + assert total_frames == first_chunk_frames + self.num_subsequent_segments * subsequent_chunk_frames, \ + f"Total frames ({total_frames}) must equal first_chunk_frames ({first_chunk_frames}) + " \ + f"num_subsequent_segments ({self.num_subsequent_segments}) * subsequent_chunk_frames ({subsequent_chunk_frames})" + + if not self.data_dir.exists(): + raise ValueError(f"Video directory not found: {self.data_dir}") + if not self.caption_dir.exists(): + raise ValueError(f"Caption directory not found: {self.caption_dir}") + + self.folders = [d for d in self.data_dir.iterdir() if d.is_dir()] + if len(self.folders) == 0: + raise ValueError(f"No subdirectories found in {self.data_dir}") + + # Optionally pre-filter folders with insufficient frames + # Note: This can be slow for large datasets due to IO operations + if self.filter_invalid_folders: + print(f"[MultiVideoConcatDataset] Pre-filtering {len(self.folders)} folders for sufficient frames...") + valid_folders = [] + skipped_folders = [] + for folder in self.folders: + if self._check_folder_has_enough_frames(folder): + valid_folders.append(folder) + else: + skipped_folders.append(folder.name) + + if len(skipped_folders) > 0: + print(f"[MultiVideoConcatDataset] Skipped {len(skipped_folders)} folders due to insufficient frames: {skipped_folders}") + + self.folders = valid_folders + + if len(self.folders) == 0: + raise ValueError(f"No folders with sufficient frames found in {self.data_dir}") + + # Setup transforms + self.resize_transform = transforms.Resize( + self.video_size, + antialias=True + ) + self.normalize = transforms.Normalize( + mean=[0.5, 0.5, 0.5], + std=[0.5, 0.5, 0.5] + ) + + # Lazy caches: avoid repeated decord.VideoReader / filesystem IO + self._video_info_cache = {} # video_path -> (total_frames, fps) + self._folder_files_cache = {} # folder_path -> list of video paths + + if decord is not None: + decord.bridge.set_bridge('torch') + + def _get_caption_folder(self, folder_name): + """Return the caption directory path for a given folder (sample).""" + return self.caption_dir / folder_name + + def _load_caption(self, video_path, folder_name): + """Load caption for a video file.""" + video_stem = video_path.stem + caption_folder = self._get_caption_folder(folder_name) + caption_path = caption_folder / f"{video_stem}.json" + + if caption_path.exists(): + try: + with open(caption_path, 'r', encoding='utf-8') as f: + caption_data = json.load(f) + return caption_data.get(self.caption_field, "") + except Exception: + return "" + return "" + + def _get_video_files_in_folder(self, folder_path): + """Get sorted video files in a folder, keeping only those with a per-video caption (cached).""" + key = str(folder_path) + if key in self._folder_files_cache: + return self._folder_files_cache[key] + + video_files = [] + for ext in self.video_extensions: + video_files.extend(list(folder_path.glob(f'*{ext}'))) + video_files.extend(list(folder_path.glob(f'*{ext.upper()}'))) + + def get_numeric_key(path): + try: + return int(path.stem) + except ValueError: + return float('inf') + + video_files.sort(key=get_numeric_key) + + folder_name = folder_path.name + filtered_videos = [] + for video_path in video_files: + caption = self._load_caption(video_path, folder_name) + if caption is not None and caption != "": + filtered_videos.append(video_path) + + self._folder_files_cache[key] = filtered_videos + return filtered_videos + + def _check_folder_has_enough_frames(self, folder_path): + """Check if a folder has enough total frames across all videos to complete all segments. + + This is a lenient check: we verify that the total available frames across all videos + is sufficient for the required segments, assuming ideal sampling. + """ + video_files = self._get_video_files_in_folder(folder_path) + if len(video_files) == 0: + return False + + # Calculate total available frames (in target_fps timebase) + total_available_frames = 0 + for video_path in video_files: + try: + total_frames, original_fps = self._get_video_info(video_path) + # Convert to target_fps timebase + available_in_target_fps = total_frames * self.target_fps / original_fps + total_available_frames += available_in_target_fps + except Exception: + # If we can't read a video, be conservative and skip this folder + return False + + # Check if we have enough frames for all segments + required_frames = self.total_frames + return total_available_frames >= required_frames + + def _get_video_info(self, video_path): + """Get video information without loading frames (cached).""" + key = str(video_path) + if key in self._video_info_cache: + return self._video_info_cache[key] + if decord is None: + info = self._get_video_info_ffprobe(video_path) + self._video_info_cache[key] = info + return info + try: + vr = decord.VideoReader(key, width=self.video_size[1], height=self.video_size[0]) + except: + vr = decord.VideoReader(key) + info = (len(vr), vr.get_avg_fps()) + self._video_info_cache[key] = info + return info + + @staticmethod + def _parse_fps(value): + if not value or value == "0/0": + return 0.0 + if "/" in value: + num, den = value.split("/", 1) + den_f = float(den) + return float(num) / den_f if den_f != 0 else 0.0 + return float(value) + + def _get_video_info_ffprobe(self, video_path): + cmd = [ + "ffprobe", + "-v", + "error", + "-select_streams", + "v:0", + "-count_frames", + "-show_entries", + "stream=nb_read_frames,nb_frames,avg_frame_rate,r_frame_rate,duration", + "-of", + "json", + str(video_path), + ] + try: + proc = subprocess.run(cmd, check=True, capture_output=True, text=True) + stream = json.loads(proc.stdout)["streams"][0] + fps = self._parse_fps(stream.get("avg_frame_rate")) or self._parse_fps(stream.get("r_frame_rate")) + frames_raw = stream.get("nb_read_frames") or stream.get("nb_frames") + if frames_raw and str(frames_raw).isdigit(): + total_frames = int(frames_raw) + else: + total_frames = int(round(float(stream.get("duration", 0.0)) * fps)) + if total_frames <= 0 or fps <= 0: + raise ValueError(f"Could not infer frames/fps from ffprobe output for {video_path}") + return total_frames, fps + except (subprocess.CalledProcessError, KeyError, IndexError, ValueError, json.JSONDecodeError) as exc: + raise RuntimeError(f"ffprobe failed to read video metadata for {video_path}: {exc}") from exc + + def _can_sample_from_position(self, total_frames, original_fps, num_frames, start_frame): + """Check if we can sample num_frames starting from start_frame. + + Returns True if we can sample num_frames without exceeding video bounds. + """ + if start_frame >= total_frames: + return False + + # Calculate sampling interval: original_fps / target_fps + sampling_interval = original_fps / self.target_fps + + # Calculate the last frame index we need (with rounding) + # We need to check if we can get num_frames frames + last_frame_needed = start_frame + (num_frames - 1) * sampling_interval + + # Account for rounding: the actual last frame index will be rounded + # So we need some margin to ensure we don't exceed bounds + return int(np.round(last_frame_needed)) < total_frames + + def _can_complete_all_segments_without_wrap(self, video_files, start_video_idx, start_frame): + """Check if from (start_video_idx, start_frame) we can sample all segments + without ever wrapping to the beginning (i.e. only use this video and later ones). + When single_video_only is True, all segments must come from the same video. + """ + vidx = start_video_idx + start = start_frame + + # First segment + total_frames, original_fps = self._get_video_info(video_files[vidx]) + if not self._can_sample_from_position( + total_frames, original_fps, self.first_chunk_frames, start + ): + return False + sampling_interval = original_fps / self.target_fps + next_start = start + self.first_chunk_frames * sampling_interval + if int(np.round(next_start)) >= total_frames: + if self.single_video_only: + return self.num_subsequent_segments == 0 + vidx += 1 + start = 0 + else: + start = int(np.round(next_start)) + if vidx >= len(video_files): + return False + + # Subsequent segments + for seg_i in range(self.num_subsequent_segments): + while vidx < len(video_files): + total_frames, original_fps = self._get_video_info(video_files[vidx]) + if self._can_sample_from_position( + total_frames, original_fps, self.subsequent_chunk_frames, start + ): + break + if self.single_video_only: + return False + vidx += 1 + start = 0 + if vidx >= len(video_files): + return False + sampling_interval = original_fps / self.target_fps + next_start = start + self.subsequent_chunk_frames * sampling_interval + if int(np.round(next_start)) >= total_frames: + if self.single_video_only and seg_i < self.num_subsequent_segments - 1: + return False + vidx += 1 + start = 0 + else: + start = int(np.round(next_start)) + return True + + def _sample_random_start(self, video_files): + """Sample a random (video_idx, start_frame) valid for the first segment, + and from which we can complete ALL segments without wrapping to the start. + Returns (video_idx, start_frame); falls back to (0, 0) if no valid start found. + """ + candidates = [] + for video_idx, video_path in enumerate(video_files): + total_frames, original_fps = self._get_video_info(video_path) + sampling_interval = original_fps / self.target_fps + last_needed = (self.first_chunk_frames - 1) * sampling_interval + max_start = int(np.floor(total_frames - 1 - last_needed)) + if max_start < 0: + continue + step = max(1, max_start // 50) + for start in range(0, max_start + 1, step): + if not self._can_sample_from_position( + total_frames, original_fps, self.first_chunk_frames, start + ): + continue + if self._can_complete_all_segments_without_wrap(video_files, video_idx, start): + candidates.append((video_idx, start)) + if candidates: + chosen = random.choice(candidates) + return chosen + return (0, 0) + + def _sample_frames_from_video(self, video_path, num_frames, start_frame=0): + """Sample frames from a video preserving original duration. + + Args: + video_path: Path to video file + num_frames: Number of frames to sample + start_frame: Starting frame index (for repeated sampling) + + Returns: + tuple: (frames_tensor, total_frames_in_video, original_fps) + """ + if decord is None: + return self._sample_frames_from_video_ffmpeg(video_path, num_frames, start_frame) + + try: + vr = decord.VideoReader(str(video_path), width=self.video_size[1], height=self.video_size[0]) + except: + vr = decord.VideoReader(str(video_path)) + + total_frames = len(vr) + original_fps = vr.get_avg_fps() + + if total_frames == 0: + raise ValueError(f"Video {video_path} has no frames") + + # Calculate frame sampling based on fps to preserve duration + # Calculate sampling interval: original_fps / target_fps + sampling_interval = original_fps / self.target_fps + + # Generate frame indices starting from start_frame + indices = [] + current_frame = float(start_frame) + for _ in range(num_frames): + frame_idx = int(np.round(current_frame)) + frame_idx = min(frame_idx, total_frames - 1) + indices.append(frame_idx) + current_frame += sampling_interval + if current_frame >= total_frames: + # If we run out of frames, pad with the last frame + remaining = num_frames - len(indices) + indices.extend([total_frames - 1] * remaining) + break + + indices = np.array(indices[:num_frames], dtype=np.int32) + + # Get video frames: shape (num_frames, height, width, 3) + video_frames = vr.get_batch(indices).numpy() + + # Convert to tensor and permute to (num_frames, 3, height, width) + video_tensor = torch.from_numpy(video_frames).permute(0, 3, 1, 2).contiguous() + + # Convert to float and normalize pixel values from [0, 255] to [0, 1] + video_tensor = video_tensor.float() / 255.0 + + # Resize if needed + if video_tensor.shape[2] != self.video_size[0] or video_tensor.shape[3] != self.video_size[1]: + resized_frames = [] + for i in range(video_tensor.shape[0]): + resized_frames.append(self.resize_transform(video_tensor[i])) + video_tensor = torch.stack(resized_frames, dim=0) + + # Apply normalization: (x - 0.5) / 0.5 -> range [-1, 1] + video_tensor = self.normalize(video_tensor) + video_tensor = video_tensor.to(torch.float16) + + return video_tensor, total_frames, original_fps + + def _sample_frames_from_video_ffmpeg(self, video_path, num_frames, start_frame=0): + total_frames, original_fps = self._get_video_info(video_path) + start_time = max(0.0, float(start_frame) / max(original_fps, 1e-6)) + height, width = self.video_size + cmd = [ + "ffmpeg", + "-v", + "error", + "-ss", + f"{start_time:.6f}", + "-i", + str(video_path), + "-vf", + f"fps={self.target_fps},scale={width}:{height}", + "-frames:v", + str(num_frames), + "-f", + "rawvideo", + "-pix_fmt", + "rgb24", + "pipe:1", + ] + try: + proc = subprocess.run(cmd, check=True, capture_output=True) + except subprocess.CalledProcessError as exc: + stderr = exc.stderr.decode("utf-8", errors="replace") if exc.stderr else "" + raise RuntimeError(f"ffmpeg failed to sample {video_path}: {stderr}") from exc + + frame_bytes = height * width * 3 + decoded_frames = len(proc.stdout) // frame_bytes + if decoded_frames <= 0: + raise RuntimeError(f"ffmpeg produced no frames for {video_path}") + + video_frames = np.frombuffer(proc.stdout[: decoded_frames * frame_bytes], dtype=np.uint8) + video_frames = video_frames.reshape(decoded_frames, height, width, 3) + video_tensor = torch.from_numpy(video_frames.copy()).permute(0, 3, 1, 2).contiguous() + video_tensor = video_tensor.float() / 255.0 + if decoded_frames < num_frames: + pad = video_tensor[-1:].repeat(num_frames - decoded_frames, 1, 1, 1) + video_tensor = torch.cat([video_tensor, pad], dim=0) + elif decoded_frames > num_frames: + video_tensor = video_tensor[:num_frames] + + video_tensor = self.normalize(video_tensor) + video_tensor = video_tensor.to(torch.float16) + return video_tensor, total_frames, original_fps + + def __len__(self): + return len(self.folders) + + @staticmethod + def _sample_failure(reason): + return False, reason + + def _try_get_item_from_folder(self, folder_idx, deterministic: bool = False): + """Try to get item from a specific folder. + + Returns (True, result) on success and (False, reason) on failure. + The reason is used by __getitem__ warnings so a long scan does not + look like a silent hang when most folders are too short. + """ + folder_path = self.folders[folder_idx] + folder_name = folder_path.name + + video_files = self._get_video_files_in_folder(folder_path) + + if len(video_files) == 0: + return self._sample_failure(f"{folder_name}: no videos with captions") + + # Collect all segments + all_segments = [] + prompts_list = [] + + try: + # Start position + if deterministic: + current_video_idx, current_start_frame = 0, 0 + else: + # Random start: don't always start from the first video, so we get more diversity + current_video_idx, current_start_frame = self._sample_random_start(video_files) + + # Sample first segment (9 frames) + video_path = video_files[current_video_idx] + prev_seg_video_idx = current_video_idx + total_frames, original_fps = self._get_video_info(video_path) + + # Ensure the current video is long enough for the first segment + while not self._can_sample_from_position( + total_frames, original_fps, self.first_chunk_frames, current_start_frame + ): + if self.single_video_only: + return self._sample_failure( + f"{folder_name}: first video is too short for the first chunk" + ) + current_video_idx += 1 + current_start_frame = 0 + if current_video_idx >= len(video_files): + return self._sample_failure( + f"{folder_name}: no video can provide the first chunk" + ) + video_path = video_files[current_video_idx] + prev_seg_video_idx = current_video_idx + total_frames, original_fps = self._get_video_info(video_path) + + # Sample first segment + segment_frames, total_frames, original_fps = self._sample_frames_from_video( + video_path, self.first_chunk_frames, current_start_frame + ) + all_segments.append(segment_frames) + + prompt = self._load_caption(video_path, folder_name) + prompts_list.append(prompt) + + # Update position for next sampling + sampling_interval = original_fps / self.target_fps + next_start_frame = current_start_frame + self.first_chunk_frames * sampling_interval + + # If we've exhausted this video, move to next + if int(np.round(next_start_frame)) >= total_frames: + if self.single_video_only: + if self.num_subsequent_segments > 0: + return self._sample_failure( + f"{folder_name}: single-video sample ends after first chunk" + ) + else: + current_video_idx += 1 + current_start_frame = 0 + else: + current_start_frame = int(np.round(next_start_frame)) + + chunks_from_current_video = 1 if current_video_idx == prev_seg_video_idx else 0 + + # Sample subsequent segments (12 frames each). No wrap: we only use videos from start onward. + for seg_idx in range(self.num_subsequent_segments): + if current_video_idx >= len(video_files): + if self.allow_padding: + break + return self._sample_failure( + f"{folder_name}: ran out of videos before all chunks were filled" + ) + + # Force virtual scene cut if max_shot_chunks reached: + # skip 1 second of video and treat the remainder as a new shot. + forced_scene_cut = False + if (self.max_chunks_per_shot > 0 + and chunks_from_current_video >= self.max_chunks_per_shot): + vp = video_files[current_video_idx] + _, ofps = self._get_video_info(vp) + current_start_frame += int(np.round(ofps)) + chunks_from_current_video = 0 + forced_scene_cut = True + + video_path = video_files[current_video_idx] + total_frames, original_fps = self._get_video_info(video_path) + + can_sample = True + while not self._can_sample_from_position( + total_frames, original_fps, self.subsequent_chunk_frames, current_start_frame + ): + if self.single_video_only: + can_sample = False + break + current_video_idx += 1 + current_start_frame = 0 + chunks_from_current_video = 0 + if current_video_idx >= len(video_files): + can_sample = False + break + video_path = video_files[current_video_idx] + total_frames, original_fps = self._get_video_info(video_path) + + if not can_sample: + if self.allow_padding: + break + return self._sample_failure( + f"{folder_name}: remaining videos are too short for the next chunk" + ) + + is_scene_cut = (current_video_idx != prev_seg_video_idx) or forced_scene_cut + + # Sample segment + segment_frames, total_frames, original_fps = self._sample_frames_from_video( + video_path, self.subsequent_chunk_frames, current_start_frame + ) + all_segments.append(segment_frames) + + prompt = self._load_caption(video_path, folder_name) + if is_scene_cut and self.scene_cut_prefix: + prompt = self.scene_cut_prefix + prompt + prompts_list.append(prompt) + + prev_seg_video_idx = current_video_idx + chunks_from_current_video += 1 + + # Update position for next sampling + sampling_interval = original_fps / self.target_fps + next_start_frame = current_start_frame + self.subsequent_chunk_frames * sampling_interval + + # If we've exhausted this video, move to next + if int(np.round(next_start_frame)) >= total_frames: + if self.single_video_only: + if seg_idx < self.num_subsequent_segments - 1: + return self._sample_failure( + f"{folder_name}: single-video sample is too short" + ) + else: + current_video_idx += 1 + current_start_frame = 0 + chunks_from_current_video = 0 + else: + current_start_frame = int(np.round(next_start_frame)) + + num_filled_segments = len(all_segments) + if num_filled_segments == 0: + num_valid_latent_frames = 0 + else: + num_valid_latent_frames = ( + self.first_chunk_latent_frames + + (num_filled_segments - 1) * self.num_frame_per_block + ) + + # Reject if below minimum latent frame threshold + if self.allow_padding and self.min_latent_frames > 0: + if num_valid_latent_frames < self.min_latent_frames: + return self._sample_failure( + f"{folder_name}: only {num_valid_latent_frames} valid latent frames, " + f"below min_latent_frames={self.min_latent_frames}" + ) + + if num_filled_segments < self.total_segments: + last_prompt = prompts_list[-1] if prompts_list else "" + prompts_list.extend([last_prompt] * (self.total_segments - num_filled_segments)) + + # Concatenate all segments: (total_frames, 3, height, width) + concatenated_video = torch.cat(all_segments, dim=0) + + # Ensure we have exactly total_frames + if concatenated_video.shape[0] != self.total_frames: + # Pad or trim if necessary + if concatenated_video.shape[0] < self.total_frames: + # Pad with last frame + last_frame = concatenated_video[-1:].repeat(self.total_frames - concatenated_video.shape[0], 1, 1, 1) + concatenated_video = torch.cat([concatenated_video, last_frame], dim=0) + else: + # Trim + concatenated_video = concatenated_video[:self.total_frames] + + result = { + 'frames': concatenated_video.permute(1, 0, 2, 3), + 'prompts': prompts_list, + 'idx': folder_idx + } + if self.return_image: + result['image'] = concatenated_video[0] + if self.allow_padding: + result['num_valid_latent_frames'] = num_valid_latent_frames + return True, result + except Exception as exc: + return self._sample_failure(f"{folder_name}: {type(exc).__name__}: {exc}") + + def __getitem__(self, idx): + start_time = time.monotonic() + last_warning_time = start_time + attempts = 0 + last_failure = None + + def maybe_warn(folder_idx, failure_reason): + nonlocal last_warning_time + if self.sample_warning_seconds <= 0: + return + elapsed = time.monotonic() - start_time + if elapsed < self.sample_warning_seconds: + return + if ( + self.sample_warning_interval_seconds > 0 + and time.monotonic() - last_warning_time < self.sample_warning_interval_seconds + ): + return + last_warning_time = time.monotonic() + folder_name = self.folders[folder_idx % len(self.folders)].name + warnings.warn( + "[MultiVideoConcatDataset] Still searching for a valid sample " + f"after {elapsed:.1f}s and {attempts} folder attempts " + f"(requested_idx={idx}, current_folder={folder_name}, " + f"last_failure={failure_reason}). This usually means the dataset " + f"does not contain enough video duration for total_frames={self.total_frames} " + f"at target_fps={self.target_fps}. Consider reducing the training " + "window, enabling allow_padding, lowering min_latent_frames, or " + "pre-filtering invalid folders.", + RuntimeWarning, + stacklevel=2, + ) + + # First try the requested folder + attempts += 1 + success, result = self._try_get_item_from_folder(idx, deterministic=self.deterministic) + if success: + return result + last_failure = result + maybe_warn(idx, last_failure) + + # If the requested folder fails, try other folders. + # If any valid folder exists in the dataset, try to return data from it: + # scan every other folder starting from idx + 1 and return the first + # successful sample. + num_folders = len(self.folders) + for i in range(1, num_folders): + alt_idx = (idx + i) % num_folders + attempts += 1 + success, result = self._try_get_item_from_folder( + alt_idx, + deterministic=self.deterministic, + ) + if success: + return result + last_failure = result + maybe_warn(alt_idx, last_failure) + + # If all attempts fail, raise an error + elapsed = time.monotonic() - start_time + if self.sample_warning_seconds > 0 and elapsed >= self.sample_warning_seconds: + warnings.warn( + "[MultiVideoConcatDataset] No valid sample was found after " + f"{elapsed:.1f}s and {attempts} folder attempts " + f"(requested_idx={idx}, last_failure={last_failure}).", + RuntimeWarning, + stacklevel=2, + ) + raise ValueError( + f"Failed to sample valid data from folder {idx} and nearby folders. " + f"This may indicate insufficient valid videos in the dataset. " + f"Tried {attempts} folders in {elapsed:.1f}s. Last failure: {last_failure}" + ) + + +def cycle(dl): + while True: + for data in dl: + yield data + +def multi_video_collate_fn(batch): + # batch is a length-B list of dictionaries returned by __getitem__. + frames = torch.stack([b["frames"] for b in batch], dim=0) # (B, T, C, H, W) + + # Keep prompts as one list per sample: + # [[p0_seg0, p0_seg1, ...], [p1_seg0, ...], ...]. + prompts_list = [b["prompts"] for b in batch] # List[List[str]] + + idx = torch.tensor([b["idx"] for b in batch], dtype=torch.long) + + result = { + "frames": frames, + "prompts": prompts_list, + "idx": idx, + } + + if "image" in batch[0]: + result["image"] = torch.stack([b["image"] for b in batch], dim=0) + + if "num_valid_latent_frames" in batch[0]: + result["num_valid_latent_frames"] = torch.tensor( + [b["num_valid_latent_frames"] for b in batch], dtype=torch.long + ) + + return result + + +def eval_collate_fn(batch): + """Collate for text-only datasets (no frames).""" + prompts_list = [b["prompts"] for b in batch] + idx = torch.tensor([b["idx"] for b in batch], dtype=torch.long) + result = { + "prompts": prompts_list, + "idx": idx, + } + if "shot_durations" in batch[0]: + result["shot_durations"] = [b["shot_durations"] for b in batch] + return result diff --git a/utils/distributed.py b/utils/distributed.py new file mode 100644 index 0000000..acd5144 --- /dev/null +++ b/utils/distributed.py @@ -0,0 +1,149 @@ +from datetime import timedelta +from functools import partial +import os +import torch +import torch.distributed as dist +from torch.distributed.fsdp import FullStateDictConfig, FullyShardedDataParallel as FSDP, MixedPrecision, ShardingStrategy, StateDictType +from torch.distributed.fsdp.api import CPUOffload +from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy + + +def fsdp_state_dict(model): + fsdp_fullstate_save_policy = FullStateDictConfig( + offload_to_cpu=True, rank0_only=True + ) + with FSDP.state_dict_type( + model, StateDictType.FULL_STATE_DICT, fsdp_fullstate_save_policy + ): + checkpoint = model.state_dict() + + return checkpoint + + +def fsdp_wrap(module, sharding_strategy="full", mixed_precision=False, wrap_strategy="size", min_num_params=int(5e7), transformer_module=None, cpu_offload=False): + if mixed_precision: + mixed_precision_policy = MixedPrecision( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + buffer_dtype=torch.float32, + cast_forward_inputs=False + ) + else: + mixed_precision_policy = None + + if wrap_strategy == "transformer": + auto_wrap_policy = partial( + transformer_auto_wrap_policy, + transformer_layer_cls=transformer_module + ) + elif wrap_strategy == "size": + auto_wrap_policy = partial( + size_based_auto_wrap_policy, + min_num_params=min_num_params + ) + else: + raise ValueError(f"Invalid wrap strategy: {wrap_strategy}") + + os.environ["NCCL_CROSS_NIC"] = "1" + + sharding_strategy = { + "full": ShardingStrategy.FULL_SHARD, + "hybrid_full": ShardingStrategy.HYBRID_SHARD, + "hybrid_zero2": ShardingStrategy._HYBRID_SHARD_ZERO2, + "no_shard": ShardingStrategy.NO_SHARD, + }[sharding_strategy] + + module = FSDP( + module, + auto_wrap_policy=auto_wrap_policy, + sharding_strategy=sharding_strategy, + mixed_precision=mixed_precision_policy, + device_id=torch.cuda.current_device(), + limit_all_gathers=True, + use_orig_params=True, + cpu_offload=CPUOffload(offload_params=cpu_offload), + sync_module_states=False # Load ckpt on rank 0 and sync to other ranks + ) + return module + + +def barrier(): + if dist.is_initialized(): + dist.barrier() + + +def launch_distributed_job(backend: str = "nccl"): + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + host = os.environ["MASTER_ADDR"] + port = int(os.environ["MASTER_PORT"]) + + if ":" in host: # IPv6 + init_method = f"tcp://[{host}]:{port}" + else: # IPv4 + init_method = f"tcp://{host}:{port}" + # Use a long timeout so that slow collectives during checkpoint saving + # (e.g. FSDP.optim_state_dict all-gather + rank0-only disk write for a + # multi-GB full optimizer state) do not trip the NCCL watchdog on other + # ranks while they wait at the post-save barrier. + dist.init_process_group(rank=rank, world_size=world_size, backend=backend, + init_method=init_method, timeout=timedelta(minutes=60)) + torch.cuda.set_device(local_rank) + + +class EMA_FSDP: + def __init__(self, fsdp_module: torch.nn.Module, decay: float = 0.999): + self.decay = decay + self.shadow = {} + self._init_shadow(fsdp_module) + + @staticmethod + def _clean_param_name(name: str) -> str: + """Remove FSDP wrapper prefixes from parameter names.""" + return name.replace("_fsdp_wrapped_module.", "").replace("_checkpoint_wrapped_module.", "").replace("_orig_mod.", "") + + @torch.no_grad() + def _init_shadow(self, fsdp_module): + from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + with FSDP.summon_full_params(fsdp_module, writeback=False): + for n, p in fsdp_module.module.named_parameters(): + # Clean the parameter name to remove FSDP prefixes + # This ensures shadow keys are compatible with unwrapped models for inference + cleaned_name = self._clean_param_name(n) + self.shadow[cleaned_name] = p.detach().clone().float().cpu() + + @torch.no_grad() + def update(self, fsdp_module): + d = self.decay + from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + with FSDP.summon_full_params(fsdp_module, writeback=False): + for n, p in fsdp_module.module.named_parameters(): + cleaned_name = self._clean_param_name(n) + if cleaned_name in self.shadow: + self.shadow[cleaned_name].mul_(d).add_(p.detach().float().cpu(), alpha=1. - d) + + # Optional helpers --------------------------------------------------- + def state_dict(self): + # Return shadow dict directly - keys are already cleaned during init/update + # This makes the state_dict directly usable for inference with unwrapped models + return self.shadow # picklable + + def load_state_dict(self, sd): + # Handle both cases: with or without FSDP prefixes + # This ensures backward compatibility and flexibility + cleaned_sd = {} + for k, v in sd.items(): + # Remove FSDP prefixes if present to match internal naming convention + cleaned_key = self._clean_param_name(k) + cleaned_sd[cleaned_key] = v.clone() + self.shadow = cleaned_sd + + def copy_to(self, fsdp_module): + # load EMA weights into an (unwrapped) copy of the generator + from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + with FSDP.summon_full_params(fsdp_module, writeback=True): + for n, p in fsdp_module.module.named_parameters(): + cleaned_name = self._clean_param_name(n) + if cleaned_name in self.shadow: + p.data.copy_(self.shadow[cleaned_name].to(p.dtype, device=p.device)) diff --git a/utils/error_buffer.py b/utils/error_buffer.py new file mode 100644 index 0000000..678c3b3 --- /dev/null +++ b/utils/error_buffer.py @@ -0,0 +1,321 @@ +# Adopted from https://github.com/vita-epfl/Stable-Video-Infinity +# SPDX-License-Identifier: Apache-2.0 + +import random +import torch + + +class ErrorBuffer: + """Bucketed ring buffer for storing prediction errors on CPU. + + Two layouts are supported: + + * **1D (timestep-only)** — when ``num_blocks <= 0``. Buckets are keyed by + the diffusion timestep. This is the original SVI behavior. + + * **2D (position × timestep)** — when ``num_blocks > 0``. Each entry is + keyed by both the global block position along the sequence and the + timestep. Inject paths can then choose: + - ``sample(pos, t)``: match BOTH position and timestep + (E_vid / E_noise — noise-level dependent errors) + - ``sample_pos_any_t(pos)``: match position, sample uniformly across + timesteps (E_img — position-dependent context corruption that is + agnostic to the current denoising step) + - ``sample_global()``: legacy fallback, samples uniformly everywhere + + The 2D layout encodes the teacher-forcing insight that ``noisy_suffix[i]`` + looks at clean_prefix[0..i] during training but at model rollouts during + inference; storing prediction errors per-position therefore lets later + blocks self-feed larger errors without any manual position ramp. + + **Sharded timestep buckets** (``shard_size > 1``): + Each rank only allocates the timestep buckets it owns + (``t_bucket % shard_size == shard_rank``), reducing per-rank CPU memory + by ~``shard_size`` times. Typically ``shard_rank/shard_size`` are set to + ``sp_rank/sp_size`` so that sharding is per-SP-rank and saving follows + the same per-SP-rank pattern as the 2D position split. On ``add()``, + non-owned buckets are silently skipped; on ``sample()``, non-owned buckets + are remapped to the nearest owned one. + """ + + def __init__( + self, + num_buckets=40, + max_size_per_bucket=50, + num_train_timesteps=1000, + modulate_factor=0.3, + replacement_strategy="random", + num_blocks=0, + global_block_offset=0, + shard_rank=0, + shard_size=1, + ): + self.num_buckets = num_buckets + self.max_size = max_size_per_bucket + self.num_train_timesteps = num_train_timesteps + self.modulate_factor = modulate_factor + self.replacement_strategy = replacement_strategy + self.bucket_width = num_train_timesteps / num_buckets + self.num_blocks = int(num_blocks) if num_blocks else 0 + # ``global_block_offset`` is only used for stats / debug display so + # users can tell which absolute positions of the full sequence this + # buffer covers (the LAST SP rank carries the highest accumulated + # error positions). It does NOT participate in bucket keying. + self.global_block_offset = int(global_block_offset) + + self.shard_rank = int(shard_rank) + self.shard_size = max(int(shard_size), 1) + self._owned_t_buckets = sorted( + t for t in range(num_buckets) if t % self.shard_size == self.shard_rank + ) + + if self.num_blocks > 0: + self.buckets = { + (p, t): [] + for p in range(self.num_blocks) + for t in self._owned_t_buckets + } + else: + self.buckets = {t: [] for t in self._owned_t_buckets} + self.total_added = 0 + + # ------------------------------------------------------------------ keys + def _t_bucket(self, timestep_index): + b = int(timestep_index / self.bucket_width) + return max(0, min(b, self.num_buckets - 1)) + + def _is_owned_t(self, t_bucket): + return self.shard_size <= 1 or (t_bucket % self.shard_size == self.shard_rank) + + def _nearest_owned_t(self, t_bucket): + """Remap ``t_bucket`` to the closest owned timestep bucket.""" + if self.shard_size <= 1 or self._is_owned_t(t_bucket): + return t_bucket + fwd = (self.shard_rank - t_bucket % self.shard_size) % self.shard_size + bwd = self.shard_size - fwd + t_up, t_down = t_bucket + fwd, t_bucket - bwd + up_ok = 0 <= t_up < self.num_buckets + down_ok = 0 <= t_down < self.num_buckets + if up_ok and down_ok: + return t_up if fwd <= bwd else t_down + return t_up if up_ok else t_down + + def _make_key(self, t_bucket, block_pos): + if self.num_blocks > 0: + assert block_pos is not None, "block_pos required when num_blocks>0" + p = max(0, min(int(block_pos), self.num_blocks - 1)) + return (p, t_bucket) + return t_bucket + + # ------------------------------------------------------------------ add + def add(self, error_block, timestep_index, block_pos=None): + """Store a single block error into the matching bucket. + + Args: + error_block: (block_size, C, H, W) tensor + timestep_index: int, raw index in [0, num_train_timesteps) + block_pos: int, global block position; required iff num_blocks>0 + """ + t = self._t_bucket(timestep_index) + if not self._is_owned_t(t): + return + key = self._make_key(t, block_pos) + # Store in the source dtype on CPU to match SVI (which keeps bf16), + # cutting buffer memory in half vs. casting to fp32. + entry = error_block.detach().to("cpu", copy=True) + + buf = self.buckets[key] + if len(buf) < self.max_size: + buf.append(entry) + else: + if self.replacement_strategy == "fifo": + buf.pop(0) + buf.append(entry) + elif self.replacement_strategy == "l2": + stacked = torch.stack(buf) + dists = (stacked - entry.unsqueeze(0)).flatten(1).norm(dim=1) + most_similar = torch.argmin(dists).item() + buf[most_similar] = entry + else: # "random" (default) + idx = random.randint(0, self.max_size - 1) + buf[idx] = entry + self.total_added += 1 + + # ------------------------------------------------------------------ sample + def sample(self, timestep_index, device, dtype, block_pos=None): + """Sample one entry matching (block_pos, timestep_index) when 2D, or + just timestep_index when 1D. Non-owned timestep buckets are + transparently remapped to the nearest owned one. Returns None if the + (remapped) bucket is empty.""" + t = self._nearest_owned_t(self._t_bucket(timestep_index)) + key = self._make_key(t, block_pos) + buf = self.buckets[key] + if not buf: + return None + err = random.choice(buf) + return self._modulate(err).to(device=device, dtype=dtype) + + def sample_pos_any_t(self, block_pos, device, dtype): + """For 2D buffers: sample at the given position, with random timestep. + + This is the natural choice for context (E_img) injection — the clean + prefix is the result of a full ODE rollout so its accumulated error + could have originated at any timestep, but its magnitude scales with + position along the sequence. + + Falls back to ``sample_global`` when the buffer is 1D. + Only owned timestep buckets are scanned. + """ + if self.num_blocks <= 0: + return self.sample_global(device, dtype) + p = max(0, min(int(block_pos), self.num_blocks - 1)) + all_entries = [] + for t in self._owned_t_buckets: + all_entries.extend(self.buckets[(p, t)]) + if not all_entries: + return None + err = random.choice(all_entries) + return self._modulate(err).to(device=device, dtype=dtype) + + def sample_global(self, device, dtype): + """Sample one entry uniformly from ALL buckets (legacy SVI E_img).""" + all_entries = [] + for buf in self.buckets.values(): + all_entries.extend(buf) + if not all_entries: + return None + err = random.choice(all_entries) + return self._modulate(err).to(device=device, dtype=dtype) + + # ------------------------------------------------------------------ misc + def _modulate(self, err): + if self.modulate_factor > 0: + lo = 1.0 - self.modulate_factor + hi = 1.0 + self.modulate_factor + err = err * random.uniform(lo, hi) + return err + + def is_empty(self): + return self.total_added == 0 + + def has_pos(self, block_pos): + """Whether ANY owned timestep bucket at ``block_pos`` has samples (2D only).""" + if self.num_blocks <= 0: + return not self.is_empty() + p = max(0, min(int(block_pos), self.num_blocks - 1)) + return any(len(self.buckets[(p, t)]) > 0 for t in self._owned_t_buckets) + + def stats(self): + filled = sum(1 for b in self.buckets.values() if len(b) > 0) + total = sum(len(b) for b in self.buckets.values()) + num_owned_t = len(self._owned_t_buckets) + denom = self.num_blocks * num_owned_t if self.num_blocks > 0 else num_owned_t + out = { + "total_added": self.total_added, + "filled_buckets": f"{filled}/{denom}", + "total_entries": total, + } + if self.shard_size > 1: + out["shard"] = f"shard_rank={self.shard_rank}/{self.shard_size} ({num_owned_t}/{self.num_buckets} t-buckets)" + if self.num_blocks > 0: + lo = self.global_block_offset + hi = self.global_block_offset + self.num_blocks + out["global_block_range"] = f"[{lo},{hi})" + return out + + def state_dict(self): + # Keys are tuples (pos, t) when 2D — torch.save handles them fine + # via pickle. We serialize the bucket layout so loaders can validate. + return { + "buckets": {k: list(v) for k, v in self.buckets.items()}, + "total_added": self.total_added, + "num_blocks": self.num_blocks, + "num_buckets": self.num_buckets, + "global_block_offset": self.global_block_offset, + "shard_rank": self.shard_rank, + "shard_size": self.shard_size, + } + + def load_state_dict(self, state, strict_offset=True): + """Restore buckets from a serialized state. + + Args: + state: dict produced by ``state_dict``. + strict_offset: when True (default) and the buffer is 2D, + refuse to load if the saved ``global_block_offset`` does + not match the current one. This prevents the silent + position-misalignment bug under SP, where a checkpoint + saved by SP rank 0 (covering global blocks ``[0, B)``) + would otherwise be loaded into SP rank 1 (which expects + ``[B, 2B)``) and corrupt position-bucketed sampling. + Pass ``strict_offset=False`` only for backward-compat + with checkpoints saved before this field existed. + """ + if self.num_blocks > 0 and strict_offset: + saved_off = state.get("global_block_offset", None) + if saved_off is None: + raise RuntimeError( + "Refusing to load: this is a 2D position-bucketed buffer " + "but the checkpoint has no `global_block_offset` field. " + "Pass strict_offset=False if you accept the misalignment risk." + ) + if int(saved_off) != self.global_block_offset: + raise RuntimeError( + f"Refusing to load: checkpoint covers global blocks " + f"starting at {saved_off}, but this rank covers blocks " + f"starting at {self.global_block_offset}. Make sure each " + f"SP rank loads its own per-rank checkpoint file." + ) + # Shard check: warn but don't crash if shard layout changed (e.g. + # resuming a non-sharded checkpoint into a sharded buffer is fine — + # we just load whichever buckets overlap). + saved_shard_size = int(state.get("shard_size", state.get("dp_size", 1))) + saved_shard_rank = int(state.get("shard_rank", state.get("dp_rank", 0))) + if saved_shard_size != self.shard_size or saved_shard_rank != self.shard_rank: + import logging + logging.warning( + f"[ErrorBuffer] Shard layout changed: checkpoint was " + f"shard_rank={saved_shard_rank}/{saved_shard_size}, current is " + f"shard_rank={self.shard_rank}/{self.shard_size}. " + f"Loading overlapping buckets only." + ) + + saved = state["buckets"] + # Lenient match: ignore keys that don't exist in the current layout. + for k in self.buckets: + if k in saved: + self.buckets[k] = saved[k] + elif isinstance(k, tuple): + # Try string-form key from older serializations + continue + elif str(k) in saved: + self.buckets[k] = saved[str(k)] + self.total_added = int(state.get("total_added", 0)) + + +def build_error_buffer(config, num_blocks=0, global_block_offset=0, + shard_rank=0, shard_size=1): + """Build an ErrorBuffer from an OmegaConf/dict config node. + + When ``num_blocks > 0`` the buffer becomes 2D (position × timestep), + enabling teacher-forcing-aware position-dependent error injection. + Pass ``global_block_offset`` so logs can identify which absolute slice + of the full sequence this rank's buffer covers (e.g. the last SP rank + is responsible for the most error-accumulated tail blocks). + + ``shard_rank`` / ``shard_size`` shard timestep buckets: each rank only + allocates the buckets it owns, reducing per-rank CPU memory by + ~``shard_size`` times. Typically set to ``(sp_rank, sp_size)``. + """ + cfg = config if isinstance(config, dict) else dict(config) + return ErrorBuffer( + num_buckets=cfg.get("num_buckets", 40), + max_size_per_bucket=cfg.get("buffer_size_per_bucket", 50), + num_train_timesteps=cfg.get("num_train_timesteps", 1000), + modulate_factor=cfg.get("modulate_factor", 0.3), + replacement_strategy=cfg.get("replacement_strategy", "random"), + num_blocks=num_blocks, + global_block_offset=global_block_offset, + shard_rank=shard_rank, + shard_size=shard_size, + ) diff --git a/utils/fp8.py b/utils/fp8.py new file mode 100644 index 0000000..d823ac3 --- /dev/null +++ b/utils/fp8.py @@ -0,0 +1,78 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# You may not use this file except in compliance with the License. +# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0 +# +# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied. +# +# SPDX-License-Identifier: Apache-2.0 +"""TorchAO FP8 post-training quantization helpers.""" + +import torch +import torch.nn as nn + + +# Small conditioning/output projections are both numerically sensitive and too +# small to amortize dynamic activation quantization on H100. +_BF16_MODULES = { + "text_embedding.0", + "text_embedding.2", + "time_embedding.0", + "time_embedding.2", + "time_projection.1", + "head.head", +} + + +def quantize_model_fp8(model: nn.Module, *, verbose: bool = False) -> int: + """Quantize compatible BF16 linear layers to row-wise dynamic FP8 in-place.""" + if not torch.cuda.is_available(): + raise RuntimeError("TorchAO FP8 inference requires a CUDA GPU.") + + device = next(model.parameters()).device + if device.type != "cuda": + raise ValueError("Move the BF16 model to CUDA before applying FP8 quantization.") + if torch.cuda.get_device_capability(device) < (8, 9): + raise RuntimeError("TorchAO FP8 inference requires compute capability 8.9 or newer.") + + try: + from torchao.quantization import ( + Float8DynamicActivationFloat8WeightConfig, + PerRow, + quantize_, + ) + except ImportError as exc: + raise ImportError( + "FP8 inference requires a TorchAO version compatible with the installed PyTorch." + ) from exc + + quantized_names = [] + skipped_names = [] + + def filter_fn(module: nn.Module, fqn: str) -> bool: + if not isinstance(module, nn.Linear): + return False + if fqn in _BF16_MODULES: + skipped_names.append(fqn) + return False + if module.weight.dtype != torch.bfloat16: + raise TypeError(f"FP8 layer {fqn!r} must be BF16, got {module.weight.dtype}.") + out_features, in_features = module.weight.shape + if in_features % 16 or out_features % 16: + skipped_names.append(fqn) + return False + quantized_names.append(fqn) + return True + + quantize_( + model, + Float8DynamicActivationFloat8WeightConfig(granularity=PerRow()), + filter_fn=filter_fn, + ) + if verbose: + print( + f"[FP8] TorchAO W8A8 quantized {len(quantized_names)} linear layers " + f"with row-wise scaling; kept {len(skipped_names)} layers in BF16" + ) + return len(quantized_names) diff --git a/utils/i2v_conditioning.py b/utils/i2v_conditioning.py new file mode 100644 index 0000000..2a12697 --- /dev/null +++ b/utils/i2v_conditioning.py @@ -0,0 +1,71 @@ +import torch + + +def _get_i2v_context_frames( + image_or_video: torch.Tensor, + initial_latent: torch.Tensor | None, +) -> int: + if initial_latent is None: + return 0 + if image_or_video.ndim != initial_latent.ndim: + raise ValueError( + f"initial_latent rank {initial_latent.ndim} must match " + f"image/video rank {image_or_video.ndim}." + ) + if image_or_video.shape[0] != initial_latent.shape[0]: + raise ValueError( + f"initial_latent batch {initial_latent.shape[0]} must match " + f"image/video batch {image_or_video.shape[0]}." + ) + if image_or_video.shape[2:] != initial_latent.shape[2:]: + raise ValueError( + f"initial_latent shape after frames {tuple(initial_latent.shape[2:])} " + f"must match image/video {tuple(image_or_video.shape[2:])}." + ) + + context_frames = int(initial_latent.shape[1]) + if context_frames <= 0: + return 0 + if context_frames >= image_or_video.shape[1]: + raise ValueError( + f"initial_latent has {context_frames} frames but clip has " + f"{image_or_video.shape[1]} frames." + ) + return context_frames + + +def _overwrite_i2v_context( + image_or_video: torch.Tensor, + initial_latent: torch.Tensor | None, + context_frames: int, +) -> torch.Tensor: + if context_frames <= 0: + return image_or_video + output = image_or_video.clone() + output[:, :context_frames] = initial_latent[:, :context_frames].to( + device=output.device, + dtype=output.dtype, + ) + return output + + +def _zero_i2v_context_timestep( + timestep: torch.Tensor, + context_frames: int, +) -> torch.Tensor: + if context_frames <= 0: + return timestep + output = timestep.clone() + output[:, :context_frames] = 0 + return output + + +def _i2v_loss_mask_like( + image_or_video: torch.Tensor, + context_frames: int, +) -> torch.Tensor | None: + if context_frames <= 0: + return None + mask = torch.ones_like(image_or_video, dtype=torch.bool) + mask[:, :context_frames] = False + return mask diff --git a/utils/inference_utils.py b/utils/inference_utils.py new file mode 100644 index 0000000..87360a4 --- /dev/null +++ b/utils/inference_utils.py @@ -0,0 +1,312 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0 +# +# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied. +# +# SPDX-License-Identifier: Apache-2.0 +"""Small helpers for release inference examples.""" + +from __future__ import annotations + +import os +from collections.abc import Mapping +from pathlib import Path +from typing import Sequence + +import torch +from einops import rearrange +from torchvision.io import write_video + +from utils.nvfp4_checkpoint import ( + clean_fsdp_state_dict_keys, + drop_fouroversix_master_weights, + is_nvfp4_state_dict, + is_te_nvfp4_checkpoint, + quantize_model_for_fouroversix_nvfp4, + quantize_model_for_transformer_engine_nvfp4, + unwrap_generator_state_dict, +) + + +def _torch_load(path: str): + try: + return torch.load(path, map_location="cpu", weights_only=False) + except TypeError: + return torch.load(path, map_location="cpu") + + +def load_generator_checkpoint(generator, checkpoint_path: str, *, use_ema: bool = False, strict: bool | None = None): + """Load a LongLive generator checkpoint into ``generator``.""" + checkpoint = _torch_load(checkpoint_path) + state_dict = unwrap_generator_state_dict(checkpoint, use_ema=use_ema) + if use_ema: + state_dict = clean_fsdp_state_dict_keys(state_dict) + if strict is None: + strict = not use_ema + return generator.load_state_dict(state_dict, strict=strict) + + +def _load_lora_state_dict(lora_ckpt_path: str) -> Mapping[str, torch.Tensor]: + """Load a LoRA checkpoint, unwrapping ``generator_lora`` when present.""" + checkpoint = _torch_load(lora_ckpt_path) + if isinstance(checkpoint, Mapping) and "generator_lora" in checkpoint: + return checkpoint["generator_lora"] + return checkpoint + + +def apply_and_merge_lora( + pipeline, + config, + *, + device: torch.device | str | None = None, + dtype: torch.dtype = torch.bfloat16, + verbose: bool = False, +): + """Wrap ``pipeline.generator.model`` with a LoRA adapter, load weights, and merge. + + The merged module ends up structurally identical to the original generator + (``nn.Linear`` layers carrying the base + LoRA delta), which is what NVFP4 + quantization needs as its starting point. + + Returns ``True`` when LoRA was applied and merged, ``False`` when the config + did not request a LoRA adapter. + """ + adapter_cfg = getattr(config, "adapter", None) + lora_ckpt = getattr(config, "lora_ckpt", None) + if adapter_cfg is None or not lora_ckpt: + return False + + import peft + from utils.lora_utils import configure_lora_for_model + + if device is not None: + pipeline.generator.to(device=torch.device(device), dtype=dtype) + else: + pipeline.generator.to(dtype=dtype) + + if verbose: + print(f"[LoRA] Wrapping generator with adapter config: {adapter_cfg}") + pipeline.generator.model = configure_lora_for_model( + pipeline.generator.model, + model_name="generator", + lora_config=adapter_cfg, + is_main_process=verbose, + ) + + if verbose: + print(f"[LoRA] Loading LoRA weights from: {lora_ckpt}") + lora_state = _load_lora_state_dict(lora_ckpt) + peft.set_peft_model_state_dict(pipeline.generator.model, lora_state) # type: ignore[arg-type] + + if verbose: + print("[LoRA] Merging LoRA delta into base weights (merge_and_unload)...") + pipeline.generator.model = pipeline.generator.model.merge_and_unload(safe_merge=True) + pipeline.generator.model.eval().requires_grad_(False) + pipeline.is_lora_enabled = False + pipeline.is_lora_merged = True + return True + + +def place_vae_for_streaming(pipeline, config) -> torch.device | None: + """Move ``pipeline.vae`` to ``config.vae_device`` for streaming-pipeline decode. + + Only acts when both ``streaming_vae`` and ``vae_device`` are set; otherwise + leaves the VAE on whatever device the rest of the pipeline already uses. + Mirrors the relocation done in ``inference.py`` so that quick-start scripts + can opt in to the streaming-pipeline VAE simply by enabling those config + fields. + """ + if not bool(getattr(config, "streaming_vae", False)): + return None + vae_device_str = getattr(config, "vae_device", None) + if not vae_device_str: + return None + + vae_device = torch.device(vae_device_str) + pipeline.vae.to(device="cpu") + pipeline.vae.to(device=vae_device) + if hasattr(pipeline.vae, "mean"): + pipeline.vae.mean = pipeline.vae.mean.to(device=vae_device) + pipeline.vae.std = pipeline.vae.std.to(device=vae_device) + return vae_device + + +def setup_nvfp4_pipeline( + pipeline, + config, + device: torch.device | str, + *, + verbose: bool = False, +): + """Configure ``pipeline`` for NVFP4 inference from a generator checkpoint. + + Handles both supported NVFP4 backends: + + * ``model_quant_use_transformer_engine=True`` -> a BF16 generator checkpoint + that gets wrapped with TransformerEngine NVFP4 modules and materialized + after moving to ``device``. + * ``model_quant_use_transformer_engine=False`` -> either a BF16 generator + checkpoint that gets quantized with FourOverSix at load time, or a + pre-materialized FourOverSix NVFP4 state dict loaded directly into the + already-quantized architecture. + + Optional LoRA support (BF16 base only): when ``config.adapter`` and + ``config.lora_ckpt`` are both set, the LoRA adapter is loaded on the BF16 + base generator, merged via ``merge_and_unload``, and the resulting weights + are then quantized — so the same yaml can swap between TE and FourOverSix + backends without pre-merging the LoRA checkpoint. + + For materialized FourOverSix checkpoints LoRA cannot be applied (the master + weights have already been quantized away); ``lora_ckpt``/``adapter`` are + ignored in that case with a printed warning. + """ + if not bool(getattr(config, "model_quant", False)): + raise ValueError("setup_nvfp4_pipeline requires model_quant=true in the config.") + + generator_ckpt = getattr(config, "generator_ckpt", None) + if not generator_ckpt: + raise ValueError("checkpoints.generator_ckpt is required for NVFP4 inference.") + + use_te = bool(getattr(config, "model_quant_use_transformer_engine", False)) + device = torch.device(device) + use_ema = bool(getattr(config, "use_ema", False)) + + checkpoint = _torch_load(generator_ckpt) + state_dict = unwrap_generator_state_dict(checkpoint, use_ema=use_ema) + if use_ema: + state_dict = clean_fsdp_state_dict_keys(state_dict) + + if is_te_nvfp4_checkpoint(checkpoint): + raise ValueError( + "Detected a TransformerEngine module state_dict export (no longer supported). " + "Re-export with `--backend transformer_engine` (merged BF16) or `--backend fouroversix`." + ) + + is_prequantized = is_nvfp4_state_dict(state_dict) + has_lora_request = bool(getattr(config, "adapter", None)) and bool(getattr(config, "lora_ckpt", None)) + + pipeline.is_lora_enabled = False + pipeline.is_lora_merged = False + + if is_prequantized: + if has_lora_request and verbose: + print( + "[NVFP4] generator_ckpt is a materialized FourOverSix NVFP4 checkpoint; " + "ignoring lora_ckpt/adapter because the master weights are already quantized. " + "Use a BF16 base checkpoint if you need to load a LoRA on top." + ) + if use_te: + raise ValueError( + "generator_ckpt is a materialized NVFP4 (FourOverSix) checkpoint; set " + "model_quant_use_transformer_engine: false." + ) + pipeline.generator.model, _ = quantize_model_for_fouroversix_nvfp4( + pipeline.generator.model, + config=config, + keep_master_weights=False, + verbose=verbose, + ) + drop_fouroversix_master_weights(pipeline.generator.model) + pipeline.generator.load_state_dict(state_dict, strict=True) + + pipeline.text_encoder.to(dtype=torch.bfloat16) + pipeline.vae.to(dtype=torch.bfloat16) + else: + load_strict = not use_ema + pipeline.generator.load_state_dict(state_dict, strict=load_strict) + + if has_lora_request: + # Apply + merge LoRA on the BF16 base before quantization. Move the + # generator to CUDA first so the TE wrapper (which requires CUDA + # modules) can later replace the merged Linear layers in-place. + apply_and_merge_lora( + pipeline, + config, + device=device, + dtype=torch.bfloat16, + verbose=verbose, + ) + + if use_te: + pipeline.generator.model, _ = quantize_model_for_transformer_engine_nvfp4( + pipeline.generator.model, + config=config, + keep_master_weights=False, + verbose=verbose, + ) + te_fallback = bool(getattr(config, "model_quant_te_fallback_to_fouroversix", False)) + if te_fallback: + from utils.quant import _materialize_mixed_quantized_weights_for_inference as materialize_fn + else: + from utils.quant import _materialize_transformer_engine_weights_for_inference as materialize_fn + else: + pipeline.generator.model, _ = quantize_model_for_fouroversix_nvfp4( + pipeline.generator.model, + config=config, + keep_master_weights=False, + verbose=verbose, + ) + from utils.quant import _materialize_quantized_weights_for_inference as materialize_fn + + pipeline.to(dtype=torch.bfloat16) + materialize_fn(pipeline.generator.model, target_device=device) + + pipeline.generator.to(device=device) + pipeline.text_encoder.to(device=device) + pipeline.vae.to(device=device) + place_vae_for_streaming(pipeline, config) + + return pipeline + + +def prepare_single_prompt_inputs( + config, + prompt: str, + device: torch.device | str, + *, + dtype: torch.dtype = torch.bfloat16, + batch_size: int = 1, + generator: torch.Generator | None = None, +): + """Create the per-block prompt list and latent noise for one text prompt.""" + num_frames = int(getattr(config, "num_output_frames", config.image_or_video_shape[1])) + frames_per_block = int(getattr(config, "num_frame_per_block", 1)) + if num_frames % frames_per_block != 0: + raise ValueError(f"num_frames={num_frames} must be divisible by num_frame_per_block={frames_per_block}") + + latent_shape = list(config.image_or_video_shape[2:]) + if len(latent_shape) != 3: + raise ValueError(f"Expected latent shape [C, H, W], got {latent_shape}") + + num_blocks = num_frames // frames_per_block + prompts = [[prompt] * num_blocks for _ in range(batch_size)] + noise = torch.randn( + [batch_size, num_frames, *latent_shape], + device=device, + dtype=dtype, + generator=generator, + ) + return noise, prompts + + +def video_to_uint8(video: torch.Tensor) -> torch.Tensor: + """Convert a generated video tensor from [T, C, H, W] or [1, T, C, H, W] to uint8 THWC.""" + if video.ndim == 5: + if video.shape[0] != 1: + raise ValueError("video_to_uint8 expects a single sample when a batch dimension is present.") + video = video[0] + if video.ndim != 4: + raise ValueError(f"Expected video tensor with 4 dims, got shape={tuple(video.shape)}") + if video.shape[1] in (1, 3): + video = rearrange(video, "t c h w -> t h w c") + return (255.0 * video.cpu()).clamp(0, 255).to(torch.uint8) + + +def save_video(video: torch.Tensor, output_path: str | os.PathLike, *, fps: int = 24) -> None: + """Save a generated LongLive video tensor as an mp4 file.""" + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + write_video(str(output_path), video_to_uint8(video), fps=fps) diff --git a/utils/kernel/README.md b/utils/kernel/README.md new file mode 100644 index 0000000..01b3039 --- /dev/null +++ b/utils/kernel/README.md @@ -0,0 +1,54 @@ + + +# LongLive KV Dequant CUDA Extension + +Build from this directory: + +```bash +cd utils/kernel +OPENBLAS_NUM_THREADS=1 OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 MAX_JOBS=4 \ + python setup.py build_ext --inplace +``` + +Runtime import: + +```python +from utils.kernel.kv_dequant import dequantize_kv_cache_fp4 +``` + +`utils.quant.dequantize_kv_cache()` already calls this extension first and falls +back to the original Triton path if the extension is not built. + +For direct calls, pass the same scale limits used by the QuantizedTensor's +`scale_rule`: + +- `static_6`: `e2m1_max=6.0`, `e4m3_max=448.0` +- `static_4`: `e2m1_max=4.0`, `e4m3_max=448.0` +- `mse` / `l1_norm` / `abs_max` 4o6 modes: `e2m1_max=6.0`, `e4m3_max=256.0` + +The normal `utils.quant.dequantize_kv_cache()` path reads these values from +`qt.scale_rule`, so manual selection is not needed there. + +You can also pass `scale_rule` directly: + +```python +out = dequantize_kv_cache_fp4( + values, + scale_factors, + amax, + num_heads=num_heads, + block_token_size=block_token_size, + dtype=torch.bfloat16, + scale_rule="static_6", +) +``` diff --git a/utils/kernel/__init__.py b/utils/kernel/__init__.py new file mode 100644 index 0000000..d0697c1 --- /dev/null +++ b/utils/kernel/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0 +# +# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied. +# +# SPDX-License-Identifier: Apache-2.0 +"""Custom CUDA kernels used by LongLive.""" diff --git a/utils/kernel/kv_dequant.cpp b/utils/kernel/kv_dequant.cpp new file mode 100644 index 0000000..8520c17 --- /dev/null +++ b/utils/kernel/kv_dequant.cpp @@ -0,0 +1,20 @@ +// Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES +// +// Licensed under the Apache License, Version 2.0 (the "License"). +// You may not use this file except in compliance with the License. +// To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0 +// +// No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied. +// +// SPDX-License-Identifier: Apache-2.0 +#include + +TORCH_LIBRARY(longlive_kernels, m) +{ + m.def("dequantize_kv_cache_fp4(Tensor[] values, Tensor[] scale_factors, Tensor[] amax, int num_heads, int block_token_size, int dtype_code, float e2m1_max, float e4m3_max) -> Tensor"); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.doc() = "LongLive custom CUDA kernels"; +} diff --git a/utils/kernel/kv_dequant.py b/utils/kernel/kv_dequant.py new file mode 100644 index 0000000..3523e6b --- /dev/null +++ b/utils/kernel/kv_dequant.py @@ -0,0 +1,82 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0 +# +# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied. +# +# SPDX-License-Identifier: Apache-2.0 +import torch + +try: + from . import longlive_kv_dequant_cuda # noqa: F401 +except ImportError: + import longlive_kv_dequant_cuda # noqa: F401 + + +def _dtype_to_code(dtype: torch.dtype) -> int: + if dtype == torch.bfloat16: + return 0 + if dtype == torch.float16: + return 1 + if dtype == torch.float32: + return 2 + raise ValueError(f"Unsupported fused KV dequant dtype: {dtype}") + + +def scale_rule_to_fp4_limits(scale_rule) -> tuple[float, float]: + """Return the dequant denominator limits used by FourOverSix ScaleRule.""" + if hasattr(scale_rule, "max_allowed_e2m1_value") and hasattr( + scale_rule, "max_allowed_e4m3_value", + ): + return ( + float(scale_rule.max_allowed_e2m1_value()), + float(scale_rule.max_allowed_e4m3_value()), + ) + + normalized = str(scale_rule).lower() + if "." in normalized: + normalized = normalized.rsplit(".", 1)[-1] + normalized = normalized.strip().strip("\"'") + + if normalized == "static_4": + return 4.0, 448.0 + if normalized == "static_6": + return 6.0, 448.0 + if normalized in {"mse", "mae", "l1_norm", "abs_max"}: + return 6.0, 256.0 + + raise ValueError(f"Unsupported FP4 scale_rule: {scale_rule}") + + +def dequantize_kv_cache_fp4( + values: list[torch.Tensor], + scale_factors: list[torch.Tensor], + amax: list[torch.Tensor], + *, + num_heads: int, + block_token_size: int, + dtype: torch.dtype, + e2m1_max: float | None = None, + e4m3_max: float | None = None, + scale_rule=None, +) -> torch.Tensor: + """Dequantize multiple AR KV-cache chunks with one CUDA launch.""" + if e2m1_max is None or e4m3_max is None: + if scale_rule is None: + raise ValueError( + "Either e2m1_max/e4m3_max or scale_rule must be provided.", + ) + e2m1_max, e4m3_max = scale_rule_to_fp4_limits(scale_rule) + + return torch.ops.longlive_kernels.dequantize_kv_cache_fp4.default( + values, + scale_factors, + amax, + num_heads, + block_token_size, + _dtype_to_code(dtype), + e2m1_max, + e4m3_max, + ) diff --git a/utils/kernel/kv_dequant_cuda.cu b/utils/kernel/kv_dequant_cuda.cu new file mode 100644 index 0000000..e4a7e7c --- /dev/null +++ b/utils/kernel/kv_dequant_cuda.cu @@ -0,0 +1,244 @@ +// Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES +// +// Licensed under the Apache License, Version 2.0 (the "License"). +// You may not use this file except in compliance with the License. +// To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0 +// +// No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied. +// +// SPDX-License-Identifier: Apache-2.0 +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace { + +#define CHECK_CUDA_TENSOR(x) TORCH_CHECK((x).is_cuda(), #x " must be a CUDA tensor") +#define CHECK_CONTIGUOUS(x) TORCH_CHECK((x).is_contiguous(), #x " must be contiguous") + +__device__ __constant__ float kE2M1ToFloat[16] = { + 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, + -0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f, +}; + +// iter-37: hardware FP4→FP16x2 via CUDA 12.8+ built-in API +// __nv_cvt_fp4x2_to_halfraw2 (wraps cvt.rn.f16x2.e2m1x2 PTX instruction). +// Returns __half2_raw with 2 fp16 values from 1 packed byte. +__device__ __forceinline__ __half2_raw e2m1x2_to_halfraw2(uint8_t byte) { + return __nv_cvt_fp4x2_to_halfraw2( + static_cast<__nv_fp4x2_storage_t>(byte), __NV_E2M1); +} + +__device__ __forceinline__ int64_t blocked_scale_index( + const int row, + const int scale_col, + const int scale_cols) +{ + // Inverse of fouroversix.quantize.utils.to_blocked for a scale matrix + // shaped [rows_padded, scale_cols]. + const int row_block = row / 128; + const int row_in_block = row - row_block * 128; + const int scale_col_block = scale_col / 4; + const int scale_col_in_block = scale_col - scale_col_block * 4; + const int scale_col_blocks = scale_cols / 4; + + const int logical_block = row_block * scale_col_blocks + scale_col_block; + return (((int64_t)logical_block * 32 + (row_in_block & 31)) * 16 + + (row_in_block >> 5) * 4 + scale_col_in_block); +} + +template +__global__ void fp4_kv_dequant_kernel( + const uint64_t* __restrict__ value_ptrs, + const uint64_t* __restrict__ scale_ptrs, + const uint64_t* __restrict__ amax_ptrs, + scalar_t* __restrict__ output, + const int64_t total_packed_values, + const int block_token_size, + const int num_heads, + const int packed_cols, + const int scale_cols, + const float inv_global_scale_denom) +{ + const int64_t packed_idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (packed_idx >= total_packed_values) { + return; + } + + const int col_pair = packed_idx % packed_cols; + const int64_t global_row = packed_idx / packed_cols; + const int rows_per_cache_block = block_token_size * num_heads; + const int cache_block = global_row / rows_per_cache_block; + const int row_in_cache_block = global_row - (int64_t)cache_block * rows_per_cache_block; + + const int token_in_block = row_in_cache_block / num_heads; + const int head = row_in_cache_block - token_in_block * num_heads; + const int out_token = cache_block * block_token_size + token_in_block; + + const auto* values = reinterpret_cast(value_ptrs[cache_block]); + const auto* scales = reinterpret_cast(scale_ptrs[cache_block]); + const auto* amax = reinterpret_cast(amax_ptrs[cache_block]); + + const uint8_t packed = values[(int64_t)row_in_cache_block * packed_cols + col_pair]; + const int scale_col = (col_pair * 2) / 16; + const int64_t scale_idx = blocked_scale_index(row_in_cache_block, scale_col, scale_cols); + const float scale = static_cast(scales[scale_idx]); + const float global_scale = amax[0] * inv_global_scale_denom; + + // iter-37: hardware FP4→FP16x2 via CUDA 12.8 built-in (wraps cvt.rn.f16x2.e2m1x2). + const __half2_raw f16x2 = e2m1x2_to_halfraw2(packed); + // __half2_raw layout: .x = low nibble's fp16 (unsigned short), .y = high nibble's. + const float low = __half2float(__ushort_as_half(f16x2.x)) * scale * global_scale; + const float high = __half2float(__ushort_as_half(f16x2.y)) * scale * global_scale; + + const int out_col = col_pair * 2; + const int64_t out_base = (((int64_t)out_token * num_heads + head) * (packed_cols * 2)) + out_col; + output[out_base] = static_cast(low); + output[out_base + 1] = static_cast(high); +} + +at::ScalarType dtype_code_to_scalar_type(const int64_t dtype_code) +{ + switch (dtype_code) { + case 0: + return at::ScalarType::BFloat16; + case 1: + return at::ScalarType::Half; + case 2: + return at::ScalarType::Float; + default: + TORCH_CHECK(false, "Unsupported KV dequant dtype code: ", dtype_code); + } + return at::ScalarType::Float; +} + +at::Tensor make_device_pointer_tensor(at::TensorList tensors) +{ + auto options = at::TensorOptions() + .dtype(at::ScalarType::Long) + .device(tensors.front().device()); + at::Tensor ptrs = at::empty({static_cast(tensors.size())}, options); + + std::vector host_ptrs(tensors.size()); + for (size_t i = 0; i < tensors.size(); ++i) { + host_ptrs[i] = reinterpret_cast(tensors[i].data_ptr()); + } + + // The pointer table is tiny; use a synchronous copy so the temporary host + // vector cannot outlive an async H2D transfer. + C10_CUDA_CHECK(cudaMemcpy( + ptrs.data_ptr(), + host_ptrs.data(), + host_ptrs.size() * sizeof(int64_t), + cudaMemcpyHostToDevice)); + return ptrs; +} + +} // namespace + +at::Tensor dequantize_kv_cache_fp4_cuda( + at::TensorList values, + at::TensorList scale_factors, + at::TensorList amax, + int64_t num_heads, + int64_t block_token_size, + int64_t dtype_code, + double e2m1_max, + double e4m3_max) +{ + TORCH_CHECK(!values.empty(), "values must contain at least one cache block"); + TORCH_CHECK(values.size() == scale_factors.size(), + "values and scale_factors must have the same length"); + TORCH_CHECK(values.size() == amax.size(), + "values and amax must have the same length"); + TORCH_CHECK(num_heads > 0, "num_heads must be positive"); + TORCH_CHECK(block_token_size > 0, "block_token_size must be positive"); + TORCH_CHECK(e2m1_max > 0.0 && e4m3_max > 0.0, + "e2m1_max and e4m3_max must be positive"); + + const auto device = values.front().device(); + c10::cuda::CUDAGuard device_guard(device); + const int64_t max_blocks = static_cast(values.size()); + const int64_t packed_cols = values.front().size(1); + const int64_t head_dim = packed_cols * 2; + const int64_t rows_padded = values.front().size(0); + const int64_t logical_rows = block_token_size * num_heads; + const int64_t scale_cols = head_dim / 16; + + TORCH_CHECK(head_dim == 128, "KV dequant currently expects head_dim=128, got ", head_dim); + TORCH_CHECK(scale_cols % 4 == 0, "scale column count must be a multiple of 4"); + TORCH_CHECK(rows_padded >= logical_rows, + "values rows are smaller than logical KV block rows"); + TORCH_CHECK(rows_padded % 128 == 0, "values rows must be padded to a multiple of 128"); + + for (int64_t i = 0; i < max_blocks; ++i) { + CHECK_CUDA_TENSOR(values[i]); + CHECK_CUDA_TENSOR(scale_factors[i]); + CHECK_CUDA_TENSOR(amax[i]); + CHECK_CONTIGUOUS(values[i]); + CHECK_CONTIGUOUS(scale_factors[i]); + CHECK_CONTIGUOUS(amax[i]); + TORCH_CHECK(values[i].device() == device, "all values tensors must be on the same device"); + TORCH_CHECK(scale_factors[i].device() == device, + "all scale_factors tensors must be on the same device"); + TORCH_CHECK(amax[i].device() == device, "all amax tensors must be on the same device"); + TORCH_CHECK(values[i].scalar_type() == at::ScalarType::Byte, + "values tensors must be uint8"); + TORCH_CHECK(amax[i].scalar_type() == at::ScalarType::Float, + "amax tensors must be float32"); + TORCH_CHECK(values[i].dim() == 2, "values tensors must be 2D"); + TORCH_CHECK(values[i].size(0) == rows_padded && values[i].size(1) == packed_cols, + "all values tensors must have the same shape"); + } + + const auto out_dtype = dtype_code_to_scalar_type(dtype_code); + at::Tensor output = at::empty( + {1, max_blocks * block_token_size, num_heads, head_dim}, + values.front().options().dtype(out_dtype)); + + cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + at::Tensor value_ptrs = make_device_pointer_tensor(values); + at::Tensor scale_ptrs = make_device_pointer_tensor(scale_factors); + at::Tensor amax_ptrs = make_device_pointer_tensor(amax); + + const int64_t total_packed_values = max_blocks * logical_rows * packed_cols; + const int threads = 256; + const dim3 blocks((total_packed_values + threads - 1) / threads); + const float inv_global_scale_denom = static_cast(1.0 / (e2m1_max * e4m3_max)); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + output.scalar_type(), + "fp4_kv_dequant_kernel", + [&] { + fp4_kv_dequant_kernel<<>>( + reinterpret_cast(value_ptrs.data_ptr()), + reinterpret_cast(scale_ptrs.data_ptr()), + reinterpret_cast(amax_ptrs.data_ptr()), + output.data_ptr(), + total_packed_values, + static_cast(block_token_size), + static_cast(num_heads), + static_cast(packed_cols), + static_cast(scale_cols), + inv_global_scale_denom); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + return output; +} + +TORCH_LIBRARY_IMPL(longlive_kernels, CUDA, m) +{ + m.impl("dequantize_kv_cache_fp4", &dequantize_kv_cache_fp4_cuda); +} diff --git a/utils/kernel/setup.py b/utils/kernel/setup.py new file mode 100644 index 0000000..3fd184b --- /dev/null +++ b/utils/kernel/setup.py @@ -0,0 +1,41 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0 +# +# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied. +# +# SPDX-License-Identifier: Apache-2.0 +from pathlib import Path + +from setuptools import setup +from torch.utils.cpp_extension import BuildExtension, CUDAExtension + + +THIS_DIR = Path(__file__).resolve().parent + +setup( + name="longlive_kv_dequant_cuda", + ext_modules=[ + CUDAExtension( + name="longlive_kv_dequant_cuda", + sources=[ + str(THIS_DIR / "kv_dequant.cpp"), + str(THIS_DIR / "kv_dequant_cuda.cu"), + ], + extra_compile_args={ + "cxx": ["-O3", "-std=c++17"], + "nvcc": [ + "-O3", + "-std=c++17", + "--expt-relaxed-constexpr", + # iter-37: need sm_100a (Blackwell arch-specific) for + # cvt.rn.f16x2.e2m1x2 instruction. Plain sm_100 lacks it. + "-gencode=arch=compute_100a,code=sm_100a", + ], + }, + ), + ], + cmdclass={"build_ext": BuildExtension}, +) diff --git a/utils/lightvae_5b_wrapper.py b/utils/lightvae_5b_wrapper.py new file mode 100644 index 0000000..7e840ca --- /dev/null +++ b/utils/lightvae_5b_wrapper.py @@ -0,0 +1,423 @@ +import logging +import os +from typing import Optional + +import torch +import torch.nn as nn + +from wan_5b.modules.vae2_2 import ( + CausalConv3d, + Decoder3d, + Encoder3d, + count_conv3d, + patchify, + unpatchify, +) + + +def _extract_checkpoint_state_dict(raw): + state = raw + if isinstance(state, dict) and "state_dict" in state: + state = state["state_dict"] + if isinstance(state, dict) and "gen_model" in state: + state = state["gen_model"] + if isinstance(state, dict) and "generator" in state: + state = state["generator"] + if not isinstance(state, dict): + raise ValueError("Unsupported checkpoint format: expected a dict-like state_dict.") + return state + + +def _map_lightvae_key_to_wanvae(key): + def _map_resnet_tail(tail): + if tail.startswith("norm1."): + return "residual.0." + tail[len("norm1."):] + if tail.startswith("conv1."): + return "residual.2." + tail[len("conv1."):] + if tail.startswith("norm2."): + return "residual.3." + tail[len("norm2."):] + if tail.startswith("conv2."): + return "residual.6." + tail[len("conv2."):] + if tail.startswith("conv_shortcut."): + return "shortcut." + tail[len("conv_shortcut."):] + return tail + + if key.startswith("dynamic_feature_projection_heads."): + return None + + if key.startswith("quant_conv."): + return key.replace("quant_conv.", "conv1.", 1) + if key.startswith("post_quant_conv."): + return key.replace("post_quant_conv.", "conv2.", 1) + + if key.startswith("encoder.conv_in."): + return key.replace("encoder.conv_in.", "encoder.conv1.", 1) + if key.startswith("encoder.mid_block.resnets.0."): + tail = key[len("encoder.mid_block.resnets.0."):] + return "encoder.middle.0." + _map_resnet_tail(tail) + if key.startswith("encoder.mid_block.attentions.0."): + return key.replace("encoder.mid_block.attentions.0.", "encoder.middle.1.", 1) + if key.startswith("encoder.mid_block.resnets.1."): + tail = key[len("encoder.mid_block.resnets.1."):] + return "encoder.middle.2." + _map_resnet_tail(tail) + if key.startswith("encoder.norm_out."): + return key.replace("encoder.norm_out.", "encoder.head.0.", 1) + if key.startswith("encoder.conv_out."): + return key.replace("encoder.conv_out.", "encoder.head.2.", 1) + + if key.startswith("encoder.down_blocks."): + parts = key.split(".") + if len(parts) >= 6 and parts[3] == "resnets": + tail = ".".join(parts[5:]) + return f"encoder.downsamples.{parts[2]}.downsamples.{parts[4]}." + _map_resnet_tail(tail) + if len(parts) >= 7 and parts[3] == "downsampler" and parts[4] == "resample": + return f"encoder.downsamples.{parts[2]}.downsamples.2.resample.{parts[5]}." + ".".join(parts[6:]) + if len(parts) >= 6 and parts[3] == "downsampler" and parts[4] == "time_conv": + return f"encoder.downsamples.{parts[2]}.downsamples.2.time_conv." + ".".join(parts[5:]) + + if key.startswith("decoder.conv_in."): + return key.replace("decoder.conv_in.", "decoder.conv1.", 1) + if key.startswith("decoder.mid_block.resnets.0."): + tail = key[len("decoder.mid_block.resnets.0."):] + return "decoder.middle.0." + _map_resnet_tail(tail) + if key.startswith("decoder.mid_block.attentions.0."): + return key.replace("decoder.mid_block.attentions.0.", "decoder.middle.1.", 1) + if key.startswith("decoder.mid_block.resnets.1."): + tail = key[len("decoder.mid_block.resnets.1."):] + return "decoder.middle.2." + _map_resnet_tail(tail) + if key.startswith("decoder.norm_out."): + return key.replace("decoder.norm_out.", "decoder.head.0.", 1) + if key.startswith("decoder.conv_out."): + return key.replace("decoder.conv_out.", "decoder.head.2.", 1) + + if key.startswith("decoder.up_blocks."): + parts = key.split(".") + if len(parts) >= 6 and parts[3] == "resnets": + tail = ".".join(parts[5:]) + return f"decoder.upsamples.{parts[2]}.upsamples.{parts[4]}." + _map_resnet_tail(tail) + if len(parts) >= 7 and parts[3] == "upsampler" and parts[4] == "resample": + return f"decoder.upsamples.{parts[2]}.upsamples.3.resample.{parts[5]}." + ".".join(parts[6:]) + if len(parts) >= 6 and parts[3] == "upsampler" and parts[4] == "time_conv": + return f"decoder.upsamples.{parts[2]}.upsamples.3.time_conv." + ".".join(parts[5:]) + + return key + + +def _normalize_vae_state_dict(raw_state): + state = _extract_checkpoint_state_dict(raw_state) + normalized = {} + for key, value in state.items(): + mapped_key = _map_lightvae_key_to_wanvae(key) + if mapped_key is None: + continue + normalized[mapped_key] = value + return normalized + + +def infer_lightvae_pruning_rate_from_ckpt(vae_path, full_decoder_conv1_out=1024): + if vae_path is None or not os.path.exists(vae_path): + return None + try: + raw_state = torch.load(vae_path, map_location="cpu") + state = _extract_checkpoint_state_dict(raw_state) + except Exception as exc: + logging.warning("Failed to load checkpoint for pruning-rate inference: %s", exc) + return None + + weight = None + if isinstance(state, dict): + if "decoder.conv_in.weight" in state: + weight = state["decoder.conv_in.weight"] + elif "decoder.conv1.weight" in state: + weight = state["decoder.conv1.weight"] + + if weight is None: + try: + normalized_state = _normalize_vae_state_dict(state) + weight = normalized_state.get("decoder.conv1.weight", None) + except Exception: + weight = None + + if weight is None or not hasattr(weight, "shape") or len(weight.shape) < 1: + return None + + student_out = int(weight.shape[0]) + if full_decoder_conv1_out <= 0: + return None + + pruning_rate = 1.0 - (float(student_out) / float(full_decoder_conv1_out)) + pruning_rate = max(0.0, min(0.99, pruning_rate)) + return round(pruning_rate, 6) + + +def convert_to_channels_last_3d(module): + for child in module.children(): + if isinstance(child, nn.Conv3d): + child.weight.data = child.weight.data.to(memory_format=torch.channels_last_3d) + else: + convert_to_channels_last_3d(child) + + +class PrunableWanVAE(nn.Module): + def __init__( + self, + dim=160, + dec_dim=256, + z_dim=48, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0, + pruning_rate=0.0, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + + dim = max(1, int(round(dim * (1.0 - pruning_rate)))) + dec_dim = max(1, int(round(dec_dim * (1.0 - pruning_rate)))) + + self.encoder = Encoder3d( + dim, + z_dim * 2, + dim_mult, + num_res_blocks, + attn_scales, + self.temperal_downsample, + dropout, + ) + self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.conv2 = CausalConv3d(z_dim, z_dim, 1) + self.decoder = Decoder3d( + dec_dim, + z_dim, + dim_mult, + num_res_blocks, + attn_scales, + self.temperal_upsample, + dropout, + ) + + def encode(self, x, scale): + self.clear_cache() + x = patchify(x, patch_size=2) + total_steps = 1 + (x.shape[2] - 1) // 4 + for step in range(total_steps): + self._enc_conv_idx = [0] + if step == 0: + out = self.encoder( + x[:, :, :1, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) + else: + out_chunk = self.encoder( + x[:, :, 1 + 4 * (step - 1):1 + 4 * step, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) + out = torch.cat([out, out_chunk], 2) + mu, _ = self.conv1(out).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view( + 1, self.z_dim, 1, 1, 1 + ) + else: + mu = (mu - scale[0]) * scale[1] + self.clear_cache() + return mu + + def decode(self, z, scale): + self.clear_cache() + if isinstance(scale[0], torch.Tensor): + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( + 1, self.z_dim, 1, 1, 1 + ) + else: + z = z / scale[1] + scale[0] + total_steps = z.shape[2] + x = self.conv2(z) + for step in range(total_steps): + self._conv_idx = [0] + if step == 0: + out = self.decoder( + x[:, :, step:step + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + first_chunk=True, + ) + else: + out_chunk = self.decoder( + x[:, :, step:step + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + ) + out = torch.cat([out, out_chunk], 2) + out = unpatchify(out, patch_size=2) + self.clear_cache() + return out + + def cached_decode(self, z, scale): + if isinstance(scale[0], torch.Tensor): + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( + 1, self.z_dim, 1, 1, 1 + ) + else: + z = z / scale[1] + scale[0] + total_steps = z.shape[2] + x = self.conv2(z) + is_first = self._feat_map[0] is None + for step in range(total_steps): + self._conv_idx = [0] + if step == 0 and is_first: + out = self.decoder( + x[:, :, step:step + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + first_chunk=True, + ) + elif step == 0: + out = self.decoder( + x[:, :, step:step + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + ) + else: + out_chunk = self.decoder( + x[:, :, step:step + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + ) + out = torch.cat([out, out_chunk], 2) + return unpatchify(out, patch_size=2) + + def clear_cache(self): + self._conv_num = count_conv3d(self.decoder) + self._conv_idx = [0] + self._feat_map = [None] * self._conv_num + self._enc_conv_num = count_conv3d(self.encoder) + self._enc_conv_idx = [0] + self._enc_feat_map = [None] * self._enc_conv_num + + +def _load_lightvae_model(pretrained_path=None, z_dim=48, dim=160, device="cpu", **kwargs): + cfg = dict( + dim=dim, + z_dim=z_dim, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0, + ) + cfg.update(**kwargs) + + with torch.device("meta"): + model = PrunableWanVAE(**cfg) + + if pretrained_path is None or not os.path.exists(pretrained_path): + raise FileNotFoundError(f"VAE checkpoint not found at {pretrained_path}") + + logging.info("loading %s", pretrained_path) + raw_state = torch.load(pretrained_path, map_location="cpu") + state_dict = _normalize_vae_state_dict(raw_state) + missing, unexpected = model.load_state_dict(state_dict, strict=False, assign=True) + logging.info( + "LightVAE checkpoint loaded with strict=False (missing=%d, unexpected=%d)", + len(missing), + len(unexpected), + ) + + convert_to_channels_last_3d(model) + return model + + +class LightVAE5BWrapper(nn.Module): + def __init__( + self, + vae_path: str, + pruning_rate: Optional[float] = None, + dtype: torch.dtype = torch.bfloat16, + device: Optional[torch.device] = None, + ): + super().__init__() + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + if pruning_rate is None: + pruning_rate = infer_lightvae_pruning_rate_from_ckpt(vae_path) + if pruning_rate is None: + pruning_rate = 0.75 + logging.warning( + "Unable to infer LightVAE pruning rate from checkpoint; fallback to 0.75." + ) + + mean = [ + -0.2289, -0.0052, -0.1323, -0.2339, -0.2799, 0.0174, 0.1838, 0.1557, + -0.1382, 0.0542, 0.2813, 0.0891, 0.1570, -0.0098, 0.0375, -0.1825, + -0.2246, -0.1207, -0.0698, 0.5109, 0.2665, -0.2108, -0.2158, 0.2502, + -0.2055, -0.0322, 0.1109, 0.1567, -0.0729, 0.0899, -0.2799, -0.1230, + -0.0313, -0.1649, 0.0117, 0.0723, -0.2839, -0.2083, -0.0520, 0.3748, + 0.0152, 0.1957, 0.1433, -0.2944, 0.3573, -0.0548, -0.1681, -0.0667, + ] + std = [ + 0.4765, 1.0364, 0.4514, 1.1677, 0.5313, 0.4990, 0.4818, 0.5013, + 0.8158, 1.0344, 0.5894, 1.0901, 0.6885, 0.6165, 0.8454, 0.4978, + 0.5759, 0.3523, 0.7135, 0.6804, 0.5833, 1.4146, 0.8986, 0.5659, + 0.7069, 0.5338, 0.4889, 0.4917, 0.4069, 0.4999, 0.6866, 0.4093, + 0.5709, 0.6065, 0.6415, 0.4944, 0.5726, 1.2042, 0.5458, 1.6887, + 0.3971, 1.0600, 0.3943, 0.5537, 0.5444, 0.4089, 0.7468, 0.7744, + ] + self.mean = torch.tensor(mean, dtype=torch.float32) + self.std = torch.tensor(std, dtype=torch.float32) + self.vae_path = os.path.abspath(vae_path) + self.pruning_rate = pruning_rate + self.device = torch.device(device) + self.dtype = dtype + + self.model = _load_lightvae_model( + pretrained_path=self.vae_path, + pruning_rate=self.pruning_rate, + ).eval().requires_grad_(False) + self.to(device=self.device, dtype=self.dtype) + + def to(self, device=None, dtype=None): + device = self.device if device is None else torch.device(device) + dtype = self.dtype if dtype is None else dtype + self.model.to(device=device, dtype=dtype) + self.mean = self.mean.to(device=device, dtype=dtype) + self.std = self.std.to(device=device, dtype=dtype) + self.device = device + self.dtype = dtype + return self + + def eval(self): + super().eval() + self.model.eval() + return self + + @torch.no_grad() + def decode_to_pixel(self, latent: torch.Tensor, use_cache: bool = False) -> torch.Tensor: + zs = latent.permute(0, 2, 1, 3, 4) + if use_cache: + assert latent.shape[0] == 1, "Batch size must be 1 when using cache" + + scale = [self.mean, 1.0 / self.std] + decode_fn = self.model.cached_decode if use_cache else self.model.decode + + output = [] + for item in zs: + output.append( + decode_fn(item.unsqueeze(0).to(device=self.device, dtype=self.dtype), scale) + .float() + .clamp_(-1, 1) + .squeeze(0) + ) + output = torch.stack(output, dim=0) + return output.permute(0, 2, 1, 3, 4) diff --git a/utils/lora_utils.py b/utils/lora_utils.py new file mode 100644 index 0000000..e01c6fa --- /dev/null +++ b/utils/lora_utils.py @@ -0,0 +1,102 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0 +# +# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied. +# +# SPDX-License-Identifier: Apache-2.0 +import torch +import peft +from peft import get_peft_model_state_dict +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +from torch.distributed.fsdp import ( + StateDictType, FullStateDictConfig +) + + +def configure_lora_for_model(transformer, model_name, lora_config, is_main_process=True, all_causal=False): + """Configure LoRA for a WanDiffusionWrapper model + + Args: + transformer: The transformer model to apply LoRA to + model_name: 'generator' or 'fake_score' + lora_config: LoRA configuration + is_main_process: Whether this is the main process (for logging) + all_causal: Whether all models use causal attention blocks + + Returns: + lora_model: The LoRA-wrapped model + """ + target_linear_modules = set() + + if model_name == 'generator': + adapter_target_modules = ['CausalWanAttentionBlock'] + elif model_name == 'fake_score': + adapter_target_modules = ['CausalWanAttentionBlock'] if all_causal else ['WanAttentionBlock'] + else: + raise ValueError(f"Invalid model name: {model_name}") + + for name, module in transformer.named_modules(): + if module.__class__.__name__ in adapter_target_modules: + for full_submodule_name, submodule in module.named_modules(prefix=name): + if isinstance(submodule, torch.nn.Linear): + target_linear_modules.add(full_submodule_name) + + target_linear_modules = list(target_linear_modules) + + if is_main_process: + print(f"LoRA target modules for {model_name}: {len(target_linear_modules)} Linear layers") + if getattr(lora_config, 'verbose', False): + for module_name in sorted(target_linear_modules): + print(f" - {module_name}") + + # Create LoRA config + adapter_type = lora_config.get('type', 'lora') + if adapter_type == 'lora': + peft_config = peft.LoraConfig( + r=lora_config.get('rank', 16), + lora_alpha=lora_config.get('alpha', None) or lora_config.get('rank', 16), + lora_dropout=lora_config.get('dropout', 0.0), + target_modules=target_linear_modules, + ) + else: + raise NotImplementedError(f'Adapter type {adapter_type} is not implemented') + + # Apply LoRA to the transformer + lora_model = peft.get_peft_model(transformer, peft_config) + + if is_main_process: + print('peft_config', peft_config) + lora_model.print_trainable_parameters() + + return lora_model + + +def gather_lora_state_dict(lora_model): + with FSDP.state_dict_type( + lora_model, + StateDictType.FULL_STATE_DICT, + FullStateDictConfig(rank0_only=True, offload_to_cpu=True) + ): + full = lora_model.state_dict() + return get_peft_model_state_dict(lora_model, state_dict=full) + + +def load_lora_checkpoint(lora_model, lora_state_dict, model_name, is_main_process=True): + """Load LoRA weights from state dict + + Args: + lora_model: The LoRA-wrapped model + lora_state_dict: LoRA state dict to load + model_name: 'generator' or 'critic' + is_main_process: Whether this is the main process (for logging) + """ + if is_main_process: + print(f"Loading LoRA {model_name} weights: {len(lora_state_dict)} keys in checkpoint") + + peft.set_peft_model_state_dict(lora_model, lora_state_dict) + + if is_main_process: + print(f"LoRA {model_name} weights loaded successfully") \ No newline at end of file diff --git a/utils/loss.py b/utils/loss.py new file mode 100644 index 0000000..a0152f9 --- /dev/null +++ b/utils/loss.py @@ -0,0 +1,98 @@ +from abc import ABC, abstractmethod +import torch + + +class DenoisingLoss(ABC): + @abstractmethod + def __call__( + self, x: torch.Tensor, x_pred: torch.Tensor, + noise: torch.Tensor, noise_pred: torch.Tensor, + alphas_cumprod: torch.Tensor, + timestep: torch.Tensor, + gradient_mask: torch.Tensor = None, + **kwargs + ) -> torch.Tensor: + """ + Base class for denoising loss. + Input: + - x: the clean data with shape [B, F, C, H, W] + - x_pred: the predicted clean data with shape [B, F, C, H, W] + - noise: the noise with shape [B, F, C, H, W] + - noise_pred: the predicted noise with shape [B, F, C, H, W] + - alphas_cumprod: the cumulative product of alphas (defining the noise schedule) with shape [T] + - timestep: the current timestep with shape [B, F] + """ + pass + + +class X0PredLoss(DenoisingLoss): + def __call__( + self, x: torch.Tensor, x_pred: torch.Tensor, + noise: torch.Tensor, noise_pred: torch.Tensor, + alphas_cumprod: torch.Tensor, + timestep: torch.Tensor, + gradient_mask: torch.Tensor = None, + **kwargs + ) -> torch.Tensor: + err = (x - x_pred) ** 2 + if gradient_mask is not None: + return err[gradient_mask].mean() + return err.mean() + + +class VPredLoss(DenoisingLoss): + def __call__( + self, x: torch.Tensor, x_pred: torch.Tensor, + noise: torch.Tensor, noise_pred: torch.Tensor, + alphas_cumprod: torch.Tensor, + timestep: torch.Tensor, + gradient_mask: torch.Tensor = None, + **kwargs + ) -> torch.Tensor: + weights = 1 / (1 - alphas_cumprod[timestep].reshape(*timestep.shape, 1, 1, 1)) + err = weights * (x - x_pred) ** 2 + if gradient_mask is not None: + return err[gradient_mask].mean() + return err.mean() + + +class NoisePredLoss(DenoisingLoss): + def __call__( + self, x: torch.Tensor, x_pred: torch.Tensor, + noise: torch.Tensor, noise_pred: torch.Tensor, + alphas_cumprod: torch.Tensor, + timestep: torch.Tensor, + gradient_mask: torch.Tensor = None, + **kwargs + ) -> torch.Tensor: + err = (noise - noise_pred) ** 2 + if gradient_mask is not None: + return err[gradient_mask].mean() + return err.mean() + + +class FlowPredLoss(DenoisingLoss): + def __call__( + self, x: torch.Tensor, x_pred: torch.Tensor, + noise: torch.Tensor, noise_pred: torch.Tensor, + alphas_cumprod: torch.Tensor, + timestep: torch.Tensor, + gradient_mask: torch.Tensor = None, + **kwargs + ) -> torch.Tensor: + err = (kwargs["flow_pred"] - (noise - x)) ** 2 + if gradient_mask is not None: + return err[gradient_mask].mean() + return err.mean() + + +NAME_TO_CLASS = { + "x0": X0PredLoss, + "v": VPredLoss, + "noise": NoisePredLoss, + "flow": FlowPredLoss +} + + +def get_denoising_loss(loss_type: str) -> DenoisingLoss: + return NAME_TO_CLASS[loss_type] diff --git a/utils/memory.py b/utils/memory.py new file mode 100644 index 0000000..bf3b562 --- /dev/null +++ b/utils/memory.py @@ -0,0 +1,146 @@ +# Copied from https://github.com/lllyasviel/FramePack/tree/main/demo_utils +# Apache-2.0 License +# By lllyasviel + +import torch + + +cpu = torch.device('cpu') +gpu = torch.device(f'cuda:{torch.cuda.current_device()}') +gpu_complete_modules = [] + + +class DynamicSwapInstaller: + @staticmethod + def _install_module(module: torch.nn.Module, **kwargs): + original_class = module.__class__ + module.__dict__['forge_backup_original_class'] = original_class + + def hacked_get_attr(self, name: str): + if '_parameters' in self.__dict__: + _parameters = self.__dict__['_parameters'] + if name in _parameters: + p = _parameters[name] + if p is None: + return None + if p.__class__ == torch.nn.Parameter: + return torch.nn.Parameter(p.to(**kwargs), requires_grad=p.requires_grad) + else: + return p.to(**kwargs) + if '_buffers' in self.__dict__: + _buffers = self.__dict__['_buffers'] + if name in _buffers: + return _buffers[name].to(**kwargs) + return super(original_class, self).__getattr__(name) + + module.__class__ = type('DynamicSwap_' + original_class.__name__, (original_class,), { + '__getattr__': hacked_get_attr, + }) + + return + + @staticmethod + def _uninstall_module(module: torch.nn.Module): + if 'forge_backup_original_class' in module.__dict__: + module.__class__ = module.__dict__.pop('forge_backup_original_class') + return + + @staticmethod + def install_model(model: torch.nn.Module, **kwargs): + for m in model.modules(): + DynamicSwapInstaller._install_module(m, **kwargs) + return + + @staticmethod + def uninstall_model(model: torch.nn.Module): + for m in model.modules(): + DynamicSwapInstaller._uninstall_module(m) + return + + +def fake_diffusers_current_device(model: torch.nn.Module, target_device: torch.device): + if hasattr(model, 'scale_shift_table'): + model.scale_shift_table.data = model.scale_shift_table.data.to(target_device) + return + + for k, p in model.named_modules(): + if hasattr(p, 'weight'): + p.to(target_device) + return + + +def get_cuda_free_memory_gb(device=None): + if device is None: + device = gpu + + memory_stats = torch.cuda.memory_stats(device) + bytes_active = memory_stats['active_bytes.all.current'] + bytes_reserved = memory_stats['reserved_bytes.all.current'] + bytes_free_cuda, _ = torch.cuda.mem_get_info(device) + bytes_inactive_reserved = bytes_reserved - bytes_active + bytes_total_available = bytes_free_cuda + bytes_inactive_reserved + return bytes_total_available / (1024 ** 3) + + + +def log_gpu_memory(stage: str, device=None, rank=0): + """Log GPU memory usage at a given training stage.""" + free_gb = get_cuda_free_memory_gb(device) + total_gb = torch.cuda.get_device_properties(device).total_memory / (1024 ** 3) + used_gb = total_gb - free_gb + print(f"[rank {rank}] [GPU Memory][{stage}] Used: {used_gb:.2f} GB | Free: {free_gb:.2f} GB | Total: {total_gb:.2f} GB") + + + + +def move_model_to_device_with_memory_preservation(model, target_device, preserved_memory_gb=0): + print(f'Moving {model.__class__.__name__} to {target_device} with preserved memory: {preserved_memory_gb} GB') + + for m in model.modules(): + if get_cuda_free_memory_gb(target_device) <= preserved_memory_gb: + torch.cuda.empty_cache() + return + + if hasattr(m, 'weight'): + m.to(device=target_device) + + model.to(device=target_device) + torch.cuda.empty_cache() + return + + +def offload_model_from_device_for_memory_preservation(model, target_device, preserved_memory_gb=0): + print(f'Offloading {model.__class__.__name__} from {target_device} to preserve memory: {preserved_memory_gb} GB') + + for m in model.modules(): + if get_cuda_free_memory_gb(target_device) >= preserved_memory_gb: + torch.cuda.empty_cache() + return + + if hasattr(m, 'weight'): + m.to(device=cpu) + + model.to(device=cpu) + torch.cuda.empty_cache() + return + + +def unload_complete_models(*args): + for m in gpu_complete_modules + list(args): + m.to(device=cpu) + print(f'Unloaded {m.__class__.__name__} as complete.') + + gpu_complete_modules.clear() + torch.cuda.empty_cache() + return + + +def load_model_as_complete(model, target_device, unload=True): + if unload: + unload_complete_models() + + model.to(device=target_device) + print(f'Loaded {model.__class__.__name__} to {target_device} as complete.') + + gpu_complete_modules.append(model) + return diff --git a/utils/misc.py b/utils/misc.py new file mode 100644 index 0000000..94cf29f --- /dev/null +++ b/utils/misc.py @@ -0,0 +1,39 @@ +import numpy as np +import random +import torch + + +def set_seed(seed: int, deterministic: bool = False): + """ + Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch`. + + Args: + seed (`int`): + The seed to set. + deterministic (`bool`, *optional*, defaults to `False`): + Whether to use deterministic algorithms where available. Can slow down training. + """ + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + if deterministic: + torch.use_deterministic_algorithms(True) + + +def merge_dict_list(dict_list): + if len(dict_list) == 1: + return dict_list[0] + + merged_dict = {} + for k, v in dict_list[0].items(): + if isinstance(v, torch.Tensor): + if v.ndim == 0: + merged_dict[k] = torch.stack([d[k] for d in dict_list], dim=0) + else: + merged_dict[k] = torch.cat([d[k] for d in dict_list], dim=0) + else: + # for non-tensor values, we just copy the value from the first item + merged_dict[k] = v + return merged_dict diff --git a/utils/nvfp4_checkpoint.py b/utils/nvfp4_checkpoint.py new file mode 100644 index 0000000..c1992ac --- /dev/null +++ b/utils/nvfp4_checkpoint.py @@ -0,0 +1,169 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0 +# +# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied. +# +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from collections.abc import Mapping + +import torch +from torch import nn + + +NVFP4_CHECKPOINT_FORMAT = "longlive_generator_nvfp4" +TE_NVFP4_CHECKPOINT_FORMAT = "longlive_generator_te_nvfp4" +NVFP4_CHECKPOINT_VERSION = 1 + + +def is_nvfp4_state_dict(state_dict: object) -> bool: + """Return True when a state dict contains materialized FourOverSix NVFP4 buffers.""" + if not isinstance(state_dict, Mapping): + return False + return any(str(key).endswith("quantized_weight_values") for key in state_dict) + + +def is_te_nvfp4_checkpoint(checkpoint: object) -> bool: + """Return True for checkpoints saved with TransformerEngine module state.""" + return ( + isinstance(checkpoint, Mapping) + and checkpoint.get("checkpoint_format") == TE_NVFP4_CHECKPOINT_FORMAT + ) + + +def unwrap_generator_state_dict(checkpoint: object, use_ema: bool = False) -> object: + """Extract the generator state dict from common LongLive checkpoint layouts.""" + if not isinstance(checkpoint, Mapping): + return checkpoint + if "generator" in checkpoint or "generator_ema" in checkpoint: + ema_key = "generator_ema" if use_ema and "generator_ema" in checkpoint else "generator" + return checkpoint[ema_key] + if "model" in checkpoint: + return checkpoint["model"] + return checkpoint + + +def clean_fsdp_state_dict_keys(state_dict: Mapping[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Remove FSDP wrapper prefixes used by some EMA checkpoints.""" + return {str(key).replace("_fsdp_wrapped_module.", ""): value for key, value in state_dict.items()} + + +def build_model_quantization_config(config, keep_master_weights: bool = False): + from utils.quant import ModelQuantizationConfig + + quant_cfg = ModelQuantizationConfig( + scale_rule=getattr(config, "model_quant_scale_rule", "static_6"), + quantize_backend=getattr(config, "model_quant_backend", None), + activation_scale_rule=getattr( + config, + "model_quant_activation_scale_rule", + getattr(config, "model_quant_scale_rule", "static_6"), + ), + weight_scale_rule=getattr(config, "model_quant_weight_scale_rule", None), + gradient_scale_rule=getattr(config, "model_quant_gradient_scale_rule", None), + ) + quant_cfg.keep_master_weights = keep_master_weights + return quant_cfg + + +def _maybe_to_dict(value): + if value is None: + return None + try: + from omegaconf import OmegaConf + + if OmegaConf.is_config(value): + value = OmegaConf.to_container(value, resolve=True) + except ImportError: + pass + return dict(value) + + +def quantize_model_for_fouroversix_nvfp4(model: nn.Module, config, *, keep_master_weights: bool = False, verbose: bool = True): + """Replace eligible modules with FourOverSix NVFP4 modules using the runtime config.""" + from utils.quant import quantize_model_with_filter + + return quantize_model_with_filter( + model, + quant_config=build_model_quantization_config(config, keep_master_weights=keep_master_weights), + filtered_modules=getattr(config, "model_quant_filtered_modules", None), + use_default_filtered_modules=getattr(config, "model_quant_use_default_filtered_modules", True), + cast_model_to_bf16=True, + materialize_for_inference=False, + use_transformer_engine=False, + verbose=verbose, + ) + + +def quantize_model_for_transformer_engine_nvfp4( + model: nn.Module, + config, + *, + keep_master_weights: bool = False, + verbose: bool = True, +): + """Replace eligible modules with TransformerEngine NVFP4 wrappers.""" + from utils.quant import quantize_model_with_filter + + use_transformer_engine = True + te_inference_only = bool(getattr(config, "model_quant_te_inference_only", use_transformer_engine)) + te_low_precision_weights = bool(getattr(config, "model_quant_te_low_precision_weights", te_inference_only)) + te_fallback_to_fouroversix = bool(getattr(config, "model_quant_te_fallback_to_fouroversix", False)) + + return quantize_model_with_filter( + model, + quant_config=build_model_quantization_config(config, keep_master_weights=keep_master_weights), + filtered_modules=getattr(config, "model_quant_filtered_modules", None), + use_default_filtered_modules=getattr(config, "model_quant_use_default_filtered_modules", True), + cast_model_to_bf16=True, + materialize_for_inference=False, + use_transformer_engine=True, + te_inference_only=te_inference_only, + te_low_precision_weights=te_low_precision_weights, + te_recipe_kwargs=_maybe_to_dict(getattr(config, "model_quant_te_recipe_kwargs", None)), + te_module_kwargs=_maybe_to_dict(getattr(config, "model_quant_te_module_kwargs", None)), + te_fallback_to_fouroversix=te_fallback_to_fouroversix, + verbose=verbose, + ) + + +def drop_fouroversix_master_weights(model: nn.Module) -> list[str]: + """Drop high-precision master weights after loading materialized NVFP4 buffers.""" + materialized_modules = [] + for module_name, module in model.named_modules(): + if not hasattr(module, "parameters_to_quantize"): + continue + + parameters_to_quantize = getattr(module, "parameters_to_quantize", ()) + if callable(parameters_to_quantize): + parameters_to_quantize = parameters_to_quantize() + if not parameters_to_quantize: + continue + + dropped_any = False + for parameter_name in parameters_to_quantize: + if isinstance(getattr(module, parameter_name, None), nn.Parameter): + module.register_parameter(parameter_name, None) + dropped_any = True + elif hasattr(module, parameter_name): + setattr(module, parameter_name, None) + dropped_any = True + + if not dropped_any: + continue + for cache_name in ("_quantized_weight", "_quantized_weight_transposed", "_quantized_weights"): + if hasattr(module, cache_name): + delattr(module, cache_name) + if hasattr(module, "config") and hasattr(module.config, "keep_master_weights"): + module.config.keep_master_weights = False + materialized_modules.append(module_name) + return materialized_modules + + +def cpu_state_dict(module: nn.Module) -> dict[str, torch.Tensor]: + """Return a detached CPU state dict suitable for torch.save.""" + return {key: value.detach().cpu() for key, value in module.state_dict().items()} diff --git a/utils/nvfp4_kernel.py b/utils/nvfp4_kernel.py new file mode 100644 index 0000000..dfec6b8 --- /dev/null +++ b/utils/nvfp4_kernel.py @@ -0,0 +1,296 @@ +#!/usr/bin/env python +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0 +# +# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied. +# +# SPDX-License-Identifier: Apache-2.0 + +"""NVFP4 Fake Quantization Triton Implementation. + +This module provides high-performance GPU implementations of NVFP4 fake quantization +operations using Triton kernels. +""" + +import torch +import triton +import triton.language as tl + +__all__ = ["fp4_dequantize", "static_blockwise_fp4_fake_quant"] + + +_TORCH_TO_TL_DTYPE = { + torch.float32: tl.float32, + torch.float: tl.float32, + torch.float16: tl.float16, + torch.half: tl.float16, + torch.bfloat16: tl.bfloat16, +} + + +def _torch_dtype_to_tl(dtype: torch.dtype): + if dtype not in _TORCH_TO_TL_DTYPE: + raise ValueError(f"Unsupported dtype for fp4 fake quantization: {dtype}") + return _TORCH_TO_TL_DTYPE[dtype] + + +@triton.jit +def fp4_dequantize_kernel( + packed_ptr, + scale_ptr, + global_scale_ptr, + output_ptr, + N, + BLOCK_SIZE: tl.constexpr, + TILE_SIZE: tl.constexpr, +): + """Dequantizes FP4 packed data using per-block scaling factors. + + Args: + packed_ptr (tl.pointer): Pointer to packed uint8 tensor (M x N//2) + scale_ptr (tl.pointer): Pointer to per-block scale tensor (M x N//BLOCK_SIZE) + output_ptr (tl.pointer): Pointer to output tensor (M x N) + global_scale_ptr (tl.pointer): Pointer to global scale tensor + N (int): Number of columns in unpacked tensor + BLOCK_SIZE (tl.constexpr): Size of each FP4 quantization block + TILE_SIZE (tl.constexpr): Size of the processing tile (in packed elements) + """ + # Get program ID for processing packed elements + pid = tl.program_id(0) + + # Calculate packed element offsets (each packed element contains 2 FP4 values) + packed_start = pid * TILE_SIZE + packed_offs = packed_start + tl.arange(0, TILE_SIZE) + + # Calculate 2D coordinates for packed data + packed_row_idx = packed_offs // (N // 2) + packed_col_idx = packed_offs % (N // 2) + + # Create mask for packed data bounds checking + packed_mask = packed_col_idx < (N // 2) + + # Load global scale + global_scale = tl.load(global_scale_ptr) + + # Load packed data + packed_data = tl.load(packed_ptr + packed_offs, mask=packed_mask, other=0) + + # Unpack packed FP4 values (uint8) to float16x2 + x_f16x2_packed = tl.inline_asm_elementwise( + asm=""" + { + .reg .b8 byte0, byte1, byte2, byte3; + mov.b32 {byte0, byte1, byte2, byte3}, $4; + cvt.rn.f16x2.e2m1x2 $0, byte0; + cvt.rn.f16x2.e2m1x2 $1, byte1; + cvt.rn.f16x2.e2m1x2 $2, byte2; + cvt.rn.f16x2.e2m1x2 $3, byte3; + } + """, + constraints="=r,=r,=r,=r,r", + args=[packed_data], + dtype=tl.uint32, + is_pure=True, + pack=4, + ) + val_low = ( + (x_f16x2_packed & 0xFFFF).cast(tl.uint16).cast(tl.float16, bitcast=True).cast(tl.float32) + ) + val_high = ( + (x_f16x2_packed >> 16).cast(tl.uint16).cast(tl.float16, bitcast=True).cast(tl.float32) + ) + + # Calculate output positions for both values + out_col_low = packed_col_idx * 2 + out_col_high = packed_col_idx * 2 + 1 + out_offs_low = packed_row_idx * N + out_col_low + out_offs_high = packed_row_idx * N + out_col_high + + # Calculate block indices for scaling + block_col_low = out_col_low // BLOCK_SIZE + block_col_high = out_col_high // BLOCK_SIZE + scale_offs_low = packed_row_idx * (N // BLOCK_SIZE) + block_col_low + scale_offs_high = packed_row_idx * (N // BLOCK_SIZE) + block_col_high + + # Load scaling factors + scale_low = tl.load(scale_ptr + scale_offs_low, mask=packed_mask & (out_col_low < N), other=1.0) + scale_high = tl.load( + scale_ptr + scale_offs_high, mask=packed_mask & (out_col_high < N), other=1.0 + ) + + # Apply scaling + result_low = val_low * scale_low.to(tl.float32) * global_scale + result_high = val_high * scale_high.to(tl.float32) * global_scale + + # Store results + out_mask_low = packed_mask & (out_col_low < N) + out_mask_high = packed_mask & (out_col_high < N) + + tl.store(output_ptr + out_offs_low, result_low, mask=out_mask_low) + tl.store(output_ptr + out_offs_high, result_high, mask=out_mask_high) + + +def fp4_dequantize( + packed_tensor: torch.Tensor, + scale_tensor: torch.Tensor, + global_scale: torch.Tensor, + block_size: int = 16, + tile_size: int = 128, + dtype: torch.dtype = torch.get_default_dtype(), +) -> torch.Tensor: + """Dequantizes FP4 packed tensor using per-block scaling factors. + + Args: + packed_tensor (torch.Tensor): Packed uint8 tensor of shape (M, N//2) + scale_tensor (torch.Tensor): Per-block scale tensor of shape (M, N//block_size) + global_scale (torch.Tensor): Global scaling factor tensor + block_size (int): Size of FP4 quantization blocks + tile_size (int): Size of processing tiles + + Returns: + torch.Tensor: Dequantized tensor of shape (M, N) + """ + packed_N = packed_tensor.shape[-1] + N = packed_N * 2 + # Create output tensor with proper shape handling + output_shape = list(packed_tensor.shape) + output_shape[-1] = N + output = torch.empty(output_shape, dtype=dtype, device=packed_tensor.device) + + # Calculate total number of elements and grid size + grid = lambda meta: (triton.cdiv(packed_tensor.numel(), meta["TILE_SIZE"]),) + + fp4_dequantize_kernel[grid]( + packed_tensor, + scale_tensor, + global_scale, + output, + N, + BLOCK_SIZE=block_size, + TILE_SIZE=tile_size, + ) + + return output + + +@triton.jit +def static_blockwise_fp4_fake_quant_kernel( + x_ptr, # [NUM_FP4_BLOCKS * BLOCK_SIZE] + y_ptr, # [NUM_FP4_BLOCKS * BLOCK_SIZE] + scale_ptr, # [NUM_FP4_BLOCKS] + NUM_FP4_BLOCKS, + BLOCK_SIZE: tl.constexpr, + OUT_DTYPE: tl.constexpr, +): + pid = tl.program_id(axis=0) + if pid >= NUM_FP4_BLOCKS: + return + + block_offset = pid * BLOCK_SIZE + idx = block_offset + tl.arange(0, BLOCK_SIZE) + + scale = tl.load(scale_ptr + pid).to(tl.float32) + + x = tl.load(x_ptr + idx).to(tl.float32) + + x_abs = tl.abs(x) + # If scale is 0, inf, or nan, use 1.0 (matching CUDA kernel behavior) + # Note: (x != x) checks if x is NaN per IEEE 754 + scale_safe = tl.where( + (scale == 0) | (scale != scale) | (tl.abs(scale) == float("inf")), # noqa: PLR0124 + 1.0, + scale, + ) + abs_scaled = x_abs / scale_safe + + # FP4 values: 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0 + q_val = tl.where( + abs_scaled <= 0.25, + 0.0, + tl.where( + abs_scaled < 0.75, + 0.5, + tl.where( + abs_scaled <= 1.25, + 1.0, + tl.where( + abs_scaled < 1.75, + 1.5, + tl.where( + abs_scaled <= 2.5, + 2.0, + tl.where( + abs_scaled < 3.5, + 3.0, + tl.where(abs_scaled <= 5.0, 4.0, 6.0), + ), + ), + ), + ), + ), + ) + + x_rescaled = q_val * scale_safe + x_quant = tl.where(x >= 0, x_rescaled, -x_rescaled) + + tl.store(y_ptr + idx, x_quant.to(OUT_DTYPE)) + + +def static_blockwise_fp4_fake_quant( + x: torch.Tensor, + amax: torch.Tensor, + global_amax: torch.Tensor | None = None, + quantize_block_scales: bool = True, + out_dtype: torch.dtype | None = None, +): + """Static blockwise FP4 fake quantization using Triton kernel. + + Args: + x: [NUM_FP4_BLOCKS, BLOCK_SIZE] on CUDA. + amax: [NUM_FP4_BLOCKS] or [NUM_FP4_BLOCKS, 1] per-block amax values. + global_amax: FP32 scalar global amax. If provided, used to compute scale_fp8_quant_amax. + quantize_block_scales: If True, quantize block scales to FP8. + out_dtype: Output dtype. Defaults to x.dtype if None. + """ + assert x.ndim == 2 + NUM_FP4_BLOCKS, BLOCK_SIZE = x.shape + + if out_dtype is None: + out_dtype = x.dtype + + amax = amax.float() # Requires to be in float32 + scale = amax / 6.0 # FP4 max representable value is 6.0 + + if quantize_block_scales: + from modelopt.torch.quantization.tensor_quant import scaled_e4m3_impl + from modelopt.torch.quantization.utils import reduce_amax + + if global_amax is None: + global_amax = reduce_amax(amax, axis=None, keepdims=False, squeeze_scalar=True) + + global_amax = global_amax.float() + scale_fp8_quant_amax = global_amax / 6.0 + scale = scaled_e4m3_impl(scale, scale_fp8_quant_amax) + + x_flat = x.contiguous().view(-1) + y_flat = torch.empty_like(x_flat, dtype=out_dtype) + scale_flat = scale.view(NUM_FP4_BLOCKS).contiguous() + + tl_out_dtype = _torch_dtype_to_tl(out_dtype) + + grid = (NUM_FP4_BLOCKS,) + + with torch.cuda.device(x.device): + static_blockwise_fp4_fake_quant_kernel[grid]( + x_flat, + y_flat, + scale_flat, + NUM_FP4_BLOCKS, + BLOCK_SIZE, + OUT_DTYPE=tl_out_dtype, + ) + + return y_flat.view_as(x) \ No newline at end of file diff --git a/utils/position_embedding_utils.py b/utils/position_embedding_utils.py new file mode 100644 index 0000000..f50c028 --- /dev/null +++ b/utils/position_embedding_utils.py @@ -0,0 +1,105 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0 +# +# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied. +# +# SPDX-License-Identifier: Apache-2.0 +"""Minimal temporal RoPE helpers used by multi-shot generation.""" + +import torch + + +def select_temporal_offset_for_sample( + temporal_offset, + sample_idx: int, + f: int, + start_frame: int = 0, +): + """Select the offset slice that applies to one sample. + + ``temporal_offset`` accepts a scalar, ``[B]`` per-sample constants, + ``[F]`` shared per-frame offsets, or ``[B, F]`` per-sample per-frame + offsets. The returned value is still interpreted by + ``compute_temporal_freqs`` so full-length and local slices both work. + """ + if temporal_offset is None: + return 0.0 + if torch.is_tensor(temporal_offset): + if temporal_offset.ndim == 0: + return temporal_offset + if temporal_offset.ndim == 1: + # Usually this is a shared [F] vector. If it is too short to cover + # the requested frame range, treat it as [B] constants. + if temporal_offset.numel() == f or temporal_offset.numel() >= start_frame + f: + return temporal_offset + return temporal_offset[sample_idx] + if temporal_offset.ndim == 2: + return temporal_offset[sample_idx] + raise ValueError( + "temporal_offset tensor must be scalar, [B], [F], or [B, F], " + f"got shape={tuple(temporal_offset.shape)}" + ) + if isinstance(temporal_offset, (list, tuple)): + if not temporal_offset: + return 0.0 + if isinstance(temporal_offset[0], (list, tuple)): + return torch.as_tensor(temporal_offset[sample_idx]) + if len(temporal_offset) == f or len(temporal_offset) >= start_frame + f: + return torch.as_tensor(temporal_offset) + return temporal_offset[sample_idx] + return temporal_offset + + +def compute_temporal_freqs( + freqs_t: torch.Tensor, + f: int, + start_frame: int, + t_scale: float, + device: torch.device, + method: str = "linear", + original_seq_len: int | None = None, + temporal_offset: float = 0.0, +) -> torch.Tensor: + """Compute linear temporal RoPE freqs with an optional multi-shot offset.""" + if method != "linear": + raise ValueError(f"Only linear temporal RoPE is supported in this release, got {method}.") + if original_seq_len is not None: + raise ValueError("original_seq_len is not used by the release linear RoPE path.") + if temporal_offset is None: + temporal_offset = 0.0 + if ( + t_scale == 1.0 + and not torch.is_tensor(temporal_offset) + and float(temporal_offset) == 0.0 + ): + return freqs_t[start_frame:start_frame + f] + + base_angles = torch.angle(freqs_t[1]).to(torch.float64) + positions = torch.arange(f, device=device, dtype=torch.float64) + start_frame + if torch.is_tensor(temporal_offset): + offset = temporal_offset.to(device=device, dtype=torch.float64) + if offset.ndim == 0: + positions = positions + offset + elif offset.ndim == 1: + if offset.numel() == f: + positions = positions + offset + elif offset.numel() >= start_frame + f: + positions = positions + offset[start_frame:start_frame + f] + else: + raise ValueError( + "temporal_offset length is too short for requested RoPE " + f"range: len={offset.numel()}, start={start_frame}, f={f}" + ) + else: + raise ValueError( + "compute_temporal_freqs expects a scalar or 1D temporal_offset " + f"after sample selection, got shape={tuple(offset.shape)}" + ) + else: + positions = positions + float(temporal_offset) + positions = positions * t_scale + angles = positions.unsqueeze(-1) * base_angles.unsqueeze(0) + return torch.polar(torch.ones_like(angles), angles) diff --git a/utils/quant.py b/utils/quant.py new file mode 100644 index 0000000..7af03d6 --- /dev/null +++ b/utils/quant.py @@ -0,0 +1,818 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0 +# +# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied. +# +# SPDX-License-Identifier: Apache-2.0 +import importlib +import inspect +from contextlib import nullcontext +from copy import deepcopy +from dataclasses import dataclass +import re +from typing import Any +import warnings + +import torch +import torch.nn as nn +from fouroversix import ( + DataType, + ModelQuantizationConfig, + QuantizationConfig, + QuantizedTensor, + RoundStyle, + ScaleRule, + quantize_model, + quantize_to_fp4, +) +from fouroversix.quantize.quantized_tensor import from_blocked + +from utils.nvfp4_kernel import fp4_dequantize + +_FUSED_KV_DEQUANT_DISABLED = False +_FUSED_KV_DEQUANT_WARNED = False + +QUANTIZATION_TYPE = { + "weight": "weight", + "activation": "activation", + "kv": "kv", +} + +DEFAULT_GENERATOR_FILTERED_MODULES = [ + "text_embedding.0", + "text_embedding.2", + "patch_embedding", + "time_projection.1", + "time_embedding.0", + "time_embedding.2", + "head.head", + "head.modulation", + "re:.*norm_k$", + "re:.*norm_q$", + "re:.*norm1$", + "re:.*norm2$", + "re:.*norm3$" +] + +DEFAULT_REAL_SCORE_FILTERED_MODULES = list(DEFAULT_GENERATOR_FILTERED_MODULES) +DEFAULT_FAKE_SCORE_FILTERED_MODULES = list(DEFAULT_GENERATOR_FILTERED_MODULES) +DEFAULT_FILTERED_MODULES = list(DEFAULT_GENERATOR_FILTERED_MODULES) + +FILTER_PROFILE_ALIASES = { + "generator": "generator", + "student": "generator", + "real_score": "real_score", + "teacher": "real_score", + "fake_score": "fake_score", + "critic": "fake_score", +} + + +@dataclass +class LongLiveQuantizationConfig(QuantizationConfig): + type: str = "weight" + + def __post_init__(self) -> None: + super().__post_init__() + + if not isinstance(self.type, str): + raise TypeError("Quantization type must be a string.") + if self.type not in QUANTIZATION_TYPE: + allowed = ", ".join(QUANTIZATION_TYPE.keys()) + raise ValueError(f"Unknown quantization type '{self.type}'. Expected one of: {allowed}.") + + self.type = QUANTIZATION_TYPE[self.type] + + +def _resolve_modules_to_not_convert( + model: nn.Module, + filtered_modules: list[str] | None, +) -> list[str]: + if not filtered_modules: + return [] + + exact_names = set() + regex_patterns = [] + for pattern in filtered_modules: + if not isinstance(pattern, str): + raise TypeError("Each filtered module pattern must be a string.") + if pattern.startswith("re:"): + regex_patterns.append(re.compile(pattern[3:])) + else: + exact_names.add(pattern) + + resolved = [] + for module_name, _ in model.named_modules(): + if not module_name: + continue + if module_name in exact_names or any( + regex.search(module_name) for regex in regex_patterns + ): + resolved.append(module_name) + + return sorted(set(resolved)) + + +def _get_default_filtered_modules(filter_profile: str | None) -> list[str]: + if filter_profile is None: + return list(DEFAULT_FILTERED_MODULES) + + normalized = FILTER_PROFILE_ALIASES.get(filter_profile, filter_profile) + if normalized == "generator": + return list(DEFAULT_GENERATOR_FILTERED_MODULES) + if normalized == "real_score": + return list(DEFAULT_REAL_SCORE_FILTERED_MODULES) + if normalized == "fake_score": + return list(DEFAULT_FAKE_SCORE_FILTERED_MODULES) + + allowed = ", ".join(sorted(FILTER_PROFILE_ALIASES)) + raise ValueError( + f"Unknown filter_profile '{filter_profile}'. Expected one of: {allowed}.", + ) + + +def _warn_for_te_config_mismatch(model_quant_config: ModelQuantizationConfig) -> None: + config_entries = [("default", model_quant_config)] + module_overrides = getattr(model_quant_config, "module_config_overrides", None) or {} + config_entries.extend(sorted(module_overrides.items())) + + mismatched_rules = [] + for module_name, module_config in config_entries: + if getattr(module_config, "dtype", DataType.nvfp4) != DataType.nvfp4: + raise NotImplementedError( + "TransformerEngine replacement currently only supports NVFP4." + ) + + for attr_name in ( + "scale_rule", + "activation_scale_rule", + "weight_scale_rule", + "gradient_scale_rule", + ): + rule = getattr(module_config, attr_name, None) + if rule is not None and rule != ScaleRule.static_6: + mismatched_rules.append(f"{module_name}:{attr_name}={rule}") + + if mismatched_rules: + preview = ", ".join(mismatched_rules[:8]) + if len(mismatched_rules) > 8: + preview += ", ..." + warnings.warn( + "TransformerEngine NVFP4 path maps to `NVFP4BlockScaling` and does not " + "replicate FourOverSix non-`static_6` scale rules exactly. " + f"Mismatched config entries: {preview}", + stacklevel=3, + ) + + +def _build_te_recipe(module_config: Any, te_recipe_kwargs: dict[str, Any] | None = None): + recipe_module = importlib.import_module("transformer_engine.common.recipe") + NVFP4BlockScaling = recipe_module.NVFP4BlockScaling + + recipe_kwargs = { + "disable_2d_quantization": not getattr(module_config, "weight_scale_2d", False), + "disable_stochastic_rounding": ( + getattr(module_config, "gradient_round_style", RoundStyle.nearest) + != RoundStyle.stochastic + ), + # FourOverSix only uses RHT in specific training paths, so keep TE conservative + # by default and let callers override via `te_recipe_kwargs`. + "disable_rht": True, + } + if te_recipe_kwargs: + recipe_kwargs.update(te_recipe_kwargs) + return NVFP4BlockScaling(**recipe_kwargs) + + +class TransformerEngineLinear(nn.Module): + """A lightweight wrapper that routes a linear layer through TransformerEngine.""" + + def __init__( + self, + module: nn.Linear, + module_name: str, + module_config: Any, + inference_only: bool = False, + low_precision_weights: bool = False, + te_recipe_kwargs: dict[str, Any] | None = None, + te_module_kwargs: dict[str, Any] | None = None, + ) -> None: + super().__init__() + + try: + te = importlib.import_module("transformer_engine.pytorch") + except ImportError as exc: + raise ImportError( + "TransformerEngine is not installed, but `use_transformer_engine=True` " + "was requested." + ) from exc + + if module.weight.device.type != "cuda": + raise ValueError( + "TransformerEngine replacement requires CUDA modules. " + f"Module `{module_name}` is on `{module.weight.device}`." + ) + + self.module_name = module_name + self.in_features = module.in_features + self.out_features = module.out_features + self.inference_only = inference_only + self.low_precision_weights = low_precision_weights + self._te = te + self._recipe = _build_te_recipe( + module_config=module_config, + te_recipe_kwargs=te_recipe_kwargs, + ) + + module_kwargs = dict(te_module_kwargs or {}) + module_kwargs.setdefault("device", module.weight.device) + module_kwargs.setdefault("params_dtype", module.weight.dtype) + module_kwargs.setdefault("name", module_name) + + fp8_model_init_fn = getattr(te, "fp8_model_init", None) + if self.low_precision_weights and fp8_model_init_fn is None: + warnings.warn( + "TransformerEngine low-precision parameter init requested, but " + "`fp8_model_init` is unavailable. Falling back to regular TE parameter " + "storage for this inference path.", + stacklevel=2, + ) + self.low_precision_weights = False + + model_init_context = ( + fp8_model_init_fn( + enabled=True, + recipe=self._recipe, + preserve_high_precision_init_val=False, + ) + if self.low_precision_weights + else nullcontext() + ) + with model_init_context: + self.linear = te.Linear( + module.in_features, + module.out_features, + bias=module.bias is not None, + **module_kwargs, + ) + + self._load_from_linear(module) + + if self.inference_only: + self.linear.requires_grad_(False) + self.train(False) + else: + self.linear.weight.requires_grad_(module.weight.requires_grad) + if self.linear.bias is not None and module.bias is not None: + self.linear.bias.requires_grad_(module.bias.requires_grad) + self.train(module.training) + + def _copy_tensor_into_parameter( + self, + destination: torch.Tensor, + source: torch.Tensor, + ) -> None: + source = source.detach().to(device=destination.device) + try: + destination.copy_(source) + return + except Exception: + pass + + destination.copy_(source.to(dtype=destination.dtype)) + + def _load_from_linear(self, module: nn.Linear) -> None: + with torch.no_grad(): + try: + self._copy_tensor_into_parameter(self.linear.weight, module.weight) + if module.bias is not None and self.linear.bias is not None: + self._copy_tensor_into_parameter(self.linear.bias, module.bias) + return + except Exception as copy_exc: + state_dict = { + "weight": module.weight.detach().to(device=self.linear.weight.device), + } + if module.bias is not None: + state_dict["bias"] = module.bias.detach().to( + device=self.linear.weight.device, + ) + incompatible_keys = self.linear.load_state_dict(state_dict, strict=False) + missing_keys = [ + key + for key in getattr(incompatible_keys, "missing_keys", []) + if key != "_extra_state" + ] + unexpected_keys = list(getattr(incompatible_keys, "unexpected_keys", [])) + if missing_keys or unexpected_keys: + raise RuntimeError( + "Failed to load weights into TransformerEngine linear " + f"`{self.module_name}`. missing_keys={missing_keys}, " + f"unexpected_keys={unexpected_keys}" + ) from copy_exc + + @property + def weight(self) -> torch.Tensor: + return self.linear.weight + + @property + def bias(self) -> torch.Tensor | None: + return self.linear.bias + + def _autocast_context(self): + autocast_fn = getattr(self._te, "autocast", None) + if autocast_fn is not None: + return autocast_fn(enabled=True, recipe=self._recipe) + fp8_autocast_fn = getattr(self._te, "fp8_autocast", None) + if fp8_autocast_fn is None: + raise AttributeError( + "TransformerEngine does not expose `autocast` or `fp8_autocast`." + ) + return fp8_autocast_fn(enabled=True, fp8_recipe=self._recipe) + + def forward(self, input: torch.Tensor) -> torch.Tensor: + with self._autocast_context(): + return self.linear(input) + + def extra_repr(self) -> str: + return ( + f"in_features={self.in_features}, " + f"out_features={self.out_features}, " + f"bias={self.bias is not None}, " + f"inference_only={self.inference_only}, " + f"low_precision_weights={self.low_precision_weights}, " + "backend=transformer_engine" + ) + + +def quantize_model_with_optional_te( + model: nn.Module, + model_quant_config: ModelQuantizationConfig, + *, + use_transformer_engine: bool = False, + te_inference_only: bool = False, + te_low_precision_weights: bool | None = None, + te_recipe_kwargs: dict[str, Any] | None = None, + te_module_kwargs: dict[str, Any] | None = None, + te_fallback_to_fouroversix: bool = False, + **kwargs, +) -> list[str]: + """ + Quantize a model with FourOverSix by default, or replace `nn.Linear` with + TransformerEngine wrappers when `use_transformer_engine=True`. + """ + if not use_transformer_engine: + quantize_model(model, model_quant_config, **kwargs) + return [] + + if te_low_precision_weights is None: + te_low_precision_weights = te_inference_only + + if kwargs: + if te_fallback_to_fouroversix: + warnings.warn( + "Additional kwargs passed to `quantize_model_with_optional_te` will " + "only be forwarded to the fallback FourOverSix pass after " + f"TransformerEngine replacement: {sorted(kwargs)}", + stacklevel=2, + ) + else: + warnings.warn( + "Additional kwargs passed to `quantize_model` are ignored in the " + f"TransformerEngine path: {sorted(kwargs)}", + stacklevel=2, + ) + + _warn_for_te_config_mismatch(model_quant_config) + + replaced_modules = [] + for module_name, module in list(model.named_modules()): + if ( + module_name == "" + or module_name in model_quant_config.modules_to_not_convert + or not isinstance(module, nn.Linear) + ): + continue + + model.set_submodule( + module_name, + TransformerEngineLinear( + module=module, + module_name=module_name, + module_config=model_quant_config.get_module_config(module_name), + inference_only=te_inference_only, + low_precision_weights=te_low_precision_weights, + te_recipe_kwargs=te_recipe_kwargs, + te_module_kwargs=te_module_kwargs, + ), + ) + replaced_modules.append(module_name) + + if te_fallback_to_fouroversix: + quantize_model(model, model_quant_config, **kwargs) + + return replaced_modules + + +def _tensor_nbytes(tensor: torch.Tensor | None) -> int: + if tensor is None: + return 0 + return tensor.numel() * tensor.element_size() + + +def _materialize_transformer_engine_weights_for_inference( + model: nn.Module, + target_device: torch.device | str | None = None, + cache_transposed_weights: bool = False, +) -> tuple[list[str], int, int]: + del cache_transposed_weights + + materialized_modules = [] + master_weight_bytes = 0 + quantized_weight_bytes = 0 + + for module_name, module in model.named_modules(): + if not isinstance(module, TransformerEngineLinear): + continue + + if target_device is not None: + module.to(device=torch.device(target_device)) + + quantized_weight_bytes += _tensor_nbytes(module.weight) + quantized_weight_bytes += _tensor_nbytes(module.bias) + materialized_modules.append(module_name) + + return materialized_modules, master_weight_bytes, quantized_weight_bytes + + +def _materialize_mixed_quantized_weights_for_inference( + model: nn.Module, + target_device: torch.device | str | None = None, + cache_transposed_weights: bool = False, +) -> tuple[list[str], int, int]: + te_modules, te_master_bytes, te_quantized_bytes = ( + _materialize_transformer_engine_weights_for_inference( + model, + target_device=target_device, + cache_transposed_weights=cache_transposed_weights, + ) + ) + f46_modules, f46_master_bytes, f46_quantized_bytes = ( + _materialize_quantized_weights_for_inference( + model, + target_device=target_device, + cache_transposed_weights=cache_transposed_weights, + ) + ) + + return ( + sorted(set(te_modules + f46_modules)), + te_master_bytes + f46_master_bytes, + te_quantized_bytes + f46_quantized_bytes, + ) + + +def _materialize_quantized_weights_for_inference( + model: nn.Module, + target_device: torch.device | str | None = None, + cache_transposed_weights: bool = False, +) -> tuple[list[str], int, int]: + """ + Materialize quantized weights and drop master weights. + + Optionally cache an additional transposed quantized layout for training paths that + still require dgrad after the master weight is deleted (e.g. NVFP4 + LoRA). + + This function expects modules replaced by `fouroversix.quantize_model`. + """ + materialized_modules = [] + master_weight_bytes = 0 + quantized_weight_bytes = 0 + + for module_name, module in model.named_modules(): + if not hasattr(module, "parameters_to_quantize") or not hasattr( + module, "get_quantized_parameters", + ): + continue + + parameters_to_quantize = getattr(module, "parameters_to_quantize", ()) + if callable(parameters_to_quantize): + parameters_to_quantize = parameters_to_quantize() + if not parameters_to_quantize: + continue + + did_materialize = False + for parameter_name in parameters_to_quantize: + parameter = getattr(module, parameter_name, None) + if parameter is None: + continue + + if isinstance(parameter, nn.Parameter): + parameter_tensor = parameter.data + elif isinstance(parameter, torch.Tensor): + parameter_tensor = parameter + else: + continue + + master_weight_bytes += parameter_tensor.numel() * parameter_tensor.element_size() + get_quantized_parameters = module.get_quantized_parameters + if ( + cache_transposed_weights + and "include_transposed" in inspect.signature( + get_quantized_parameters, + ).parameters + ): + quantized_params = get_quantized_parameters( + parameter_name, + parameter_tensor, + include_transposed=True, + ) + else: + quantized_params = get_quantized_parameters( + parameter_name, + parameter_tensor, + ) + + for quantized_name, quantized_tensor in quantized_params.items(): + if not isinstance(quantized_tensor, torch.Tensor): + continue + + existing = getattr(module, quantized_name, None) + dst_dtype = ( + existing.dtype + if isinstance(existing, torch.Tensor) + else quantized_tensor.dtype + ) + if target_device is not None: + dst_device = torch.device(target_device) + elif isinstance(existing, torch.Tensor): + dst_device = existing.device + else: + dst_device = quantized_tensor.device + + quantized_tensor = quantized_tensor.to( + device=dst_device, + dtype=dst_dtype, + ) + setattr(module, quantized_name, quantized_tensor) + quantized_weight_bytes += ( + quantized_tensor.numel() * quantized_tensor.element_size() + ) + + # Drop high-precision master weight once quantized weights are materialized. + if isinstance(getattr(module, parameter_name, None), nn.Parameter): + module.register_parameter(parameter_name, None) + else: + setattr(module, parameter_name, None) + did_materialize = True + + if did_materialize: + if hasattr(module, "_quantized_weight"): + delattr(module, "_quantized_weight") + if hasattr(module, "_quantized_weight_transposed"): + delattr(module, "_quantized_weight_transposed") + if hasattr(module, "_quantized_weights"): + delattr(module, "_quantized_weights") + if hasattr(module, "config") and hasattr(module.config, "keep_master_weights"): + module.config.keep_master_weights = False + materialized_modules.append(module_name) + + return materialized_modules, master_weight_bytes, quantized_weight_bytes + + +def quantize_model_with_filter( + model: nn.Module, + quant_config: ModelQuantizationConfig | dict | None = None, + filtered_modules: list[str] | None = None, + filter_profile: str | None = None, + use_default_filtered_modules: bool = False, + cast_model_to_bf16: bool = True, + materialize_for_inference: bool = False, + materialize_target_device: torch.device | str | None = None, + use_transformer_engine: bool = False, + te_inference_only: bool = False, + te_low_precision_weights: bool | None = None, + te_recipe_kwargs: dict[str, Any] | None = None, + te_module_kwargs: dict[str, Any] | None = None, + te_fallback_to_fouroversix: bool = False, + verbose: bool = True, + **kwargs, +) -> tuple[nn.Module, list[str]]: + """ + Quantize model with FourOverSix and optionally skip selected modules. + + `filtered_modules` supports: + - Exact module names, e.g. "head.head" + - Regex patterns prefixed with "re:", e.g. "re:.*norm1$" + + `filter_profile` selects which built-in filtered module profile to use when + `use_default_filtered_modules=True`. Supported values: + "generator"/"student" and "real_score"/"teacher". + """ + if quant_config is None: + model_quant_config = ModelQuantizationConfig() + elif isinstance(quant_config, dict): + model_quant_config = ModelQuantizationConfig(**quant_config) + elif isinstance(quant_config, ModelQuantizationConfig): + model_quant_config = deepcopy(quant_config) + else: + raise TypeError( + "quant_config must be ModelQuantizationConfig, dict, or None.", + ) + + patterns = list(filtered_modules or []) + if use_default_filtered_modules: + patterns = _get_default_filtered_modules(filter_profile) + patterns + + matched_modules = _resolve_modules_to_not_convert(model, patterns) + modules_to_not_convert = set(model_quant_config.modules_to_not_convert or []) + modules_to_not_convert.update(matched_modules) + model_quant_config.modules_to_not_convert = sorted(modules_to_not_convert) + + if cast_model_to_bf16: + model.to(torch.bfloat16) + + resolved_te_low_precision_weights = ( + te_inference_only if te_low_precision_weights is None else te_low_precision_weights + ) + + te_replaced_modules = quantize_model_with_optional_te( + model, + model_quant_config, + use_transformer_engine=use_transformer_engine, + te_inference_only=te_inference_only, + te_low_precision_weights=resolved_te_low_precision_weights, + te_recipe_kwargs=te_recipe_kwargs, + te_module_kwargs=te_module_kwargs, + te_fallback_to_fouroversix=te_fallback_to_fouroversix, + **kwargs, + ) + + if materialize_for_inference: + materialize_fn = _materialize_quantized_weights_for_inference + if use_transformer_engine and te_fallback_to_fouroversix: + materialize_fn = _materialize_mixed_quantized_weights_for_inference + elif use_transformer_engine: + materialize_fn = _materialize_transformer_engine_weights_for_inference + + materialized_modules, master_bytes, quantized_bytes = materialize_fn( + model, + target_device=materialize_target_device, + ) + if verbose: + print( + "[quantize_model_with_filter] " + f"materialized_modules={len(materialized_modules)}, " + f"master_weight={master_bytes / (1024 ** 3):.3f} GiB, " + f"quantized_weight={quantized_bytes / (1024 ** 3):.3f} GiB", + ) + + if verbose: + profile_label = filter_profile or "default" + print( + "[quantize_model_with_filter] " + f"profile={profile_label}, " + f"matched={len(matched_modules)}, " + f"total_excluded={len(model_quant_config.modules_to_not_convert)}", + ) + if use_transformer_engine: + print( + "[quantize_model_with_filter] " + f"transformer_engine_replaced={len(te_replaced_modules)}, " + f"inference_only={te_inference_only}, " + f"low_precision_weights={resolved_te_low_precision_weights}, " + f"fallback_to_fouroversix={te_fallback_to_fouroversix}", + ) + + return model, matched_modules + + +def _dequantize_kv_cache_fused_cuda(kv_list, max_blocks, num_heads, block_token_size, dtype): + global _FUSED_KV_DEQUANT_DISABLED, _FUSED_KV_DEQUANT_WARNED + + if _FUSED_KV_DEQUANT_DISABLED or max_blocks <= 0: + return None + + first_qt = kv_list[0] + if first_qt.values.device.type != "cuda": + return None + + try: + from utils.kernel.kv_dequant import dequantize_kv_cache_fp4 + + blocks = kv_list[:max_blocks] + values = [qt.values for qt in blocks] + scale_factors = [qt.scale_factors for qt in blocks] + amax = [qt.amax for qt in blocks] + + return dequantize_kv_cache_fp4( + values, + scale_factors, + amax, + num_heads=num_heads, + block_token_size=block_token_size, + dtype=dtype, + scale_rule=first_qt.scale_rule, + ) + except Exception as exc: # pragma: no cover - exercised only when extension is stale/missing + _FUSED_KV_DEQUANT_DISABLED = True + if not _FUSED_KV_DEQUANT_WARNED: + warnings.warn( + "Fused CUDA KV-cache dequantization is unavailable; falling back to " + f"the Triton per-block path. Reason: {exc}", + stacklevel=2, + ) + _FUSED_KV_DEQUANT_WARNED = True + return None + + +def dequantize_kv_cache(kv_list, max_blocks, num_heads, block_token_size, dtype, device): + """ + Dequantize list of QuantizedTensor to a contiguous bf16 tensor. + kv_list[block_idx] -> QuantizedTensor(block_token_size * num_heads, 128) + Returns: [1, max_blocks * block_token_size, num_heads, 128] + """ + fused_result = _dequantize_kv_cache_fused_cuda( + kv_list, max_blocks, num_heads, block_token_size, dtype, + ) + if fused_result is not None: + return fused_result + + total_tokens = max_blocks * block_token_size + result = torch.zeros([1, total_tokens, num_heads, 128], dtype=dtype, device=device) + for block_idx in range(max_blocks): + t_start = block_idx * block_token_size + t_end = t_start + block_token_size + # deq = kv_list[block_idx].dequantize(dtype) + # triton fp4_dequantize + qt = kv_list[block_idx] + padded_shape = qt.padded_shape + scales_2d = from_blocked( + qt.scale_factors, + (padded_shape[0], padded_shape[1] // 16), + ) + global_scale = qt.amax / ( + qt.scale_rule.max_allowed_e2m1_value() + * qt.scale_rule.max_allowed_e4m3_value() + ) + deq = fp4_dequantize( + kv_list[block_idx].values, + scales_2d, + global_scale, + block_size=16, + dtype=dtype, + ) + result[0, t_start:t_end, :, :] = deq.view(block_token_size, num_heads, 128) + return result + +def clone_quantized_tensor(qt): + """Clone a QuantizedTensor by cloning its internal tensors.""" + return QuantizedTensor( + values=qt.values.clone(), + scale_factors=qt.scale_factors.clone(), + amax=qt.amax.clone() if qt.amax is not None else None, + dtype=qt.dtype, + original_shape=qt.original_shape, + scale_rule=qt.scale_rule, + padded_shape=qt.padded_shape, + ) + + +def copy_quantized_into(slot: QuantizedTensor, src: QuantizedTensor) -> None: + """In-place copy a QuantizedTensor's data into a pre-allocated slot. + + Keeps the slot's `values`/`scale_factors`/`amax` buffers persistent + (their addresses don't change) so cudagraph allocator does not see them + as fresh outputs that can be reused across step boundaries. Used by the + quantized KV cache rolling/insert paths. + """ + slot.values.copy_(src.values) + slot.scale_factors.copy_(src.scale_factors) + if src.amax is not None and slot.amax is not None: + slot.amax.copy_(src.amax) + + +def k_smooth(k: torch.Tensor) -> torch.Tensor: + return k - k.mean(dim=-1, keepdim=True) + +def quantize_kv(k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + B, S, H, D = k.shape + # B is always 1 + # S is the number of tokens + # H is the number of heads + # D is the dimension of the key and value + + config = QuantizationConfig(scale_rule="mse", backend="cuda") + # per head quantization + for head in range(H): + k_head = k[:, :, head, :] + v_head = v[:, :, head, :] + k_head = k_smooth(k_head) + v_head = v_head + k_head = quantize_to_fp4(k_head, config) + v_head = quantize_to_fp4(v_head, config) + k[:, :, head, :] = k_head + v[:, :, head, :] = v_head + return k, v diff --git a/utils/rope_triton.py b/utils/rope_triton.py new file mode 100644 index 0000000..687e3b4 --- /dev/null +++ b/utils/rope_triton.py @@ -0,0 +1,135 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0 +# +# SPDX-License-Identifier: Apache-2.0 +"""Triton RoPE kernel for causal_rope_apply. + +iter-42: replaces the complex × view_as_complex × view_as_real chain +(33 ms / 1.3% of profile + feeding elementwise muls) with a single Triton +kernel. Internal precision is fp32 — bf16 outputs cannot resolve any precision +loss from fp32 vs fp64 arithmetic at this stage. cos / sin lookup tables come +from the complex128 freqs split into real / imag floats up-front (one-shot per +freqs_i cache entry). +""" +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +# iter-45 (REVERTED): tried @triton.autotune over (BLOCK_N∈{4,8,16}, +# num_warps∈{2,4,8}, num_stages∈{1,2,3}) = 27 configs. Result was FLAT vs +# iter-42 fixed BLOCK_N=8: median tied (-0.1%), total +0.4% (autotune warmup +# bled into p1/p2 p90). The original BLOCK_N=8 / default warps was already +# near-optimal for the (N=24, D_half=64) shape, autotune found no better. +# Reverted to fixed config — same kernel as iter-42. +# +# iter-46: kernel now accepts FULL x[i] of shape [S_total, N, D] and a +# runtime `seq_len` — for rows s < seq_len it applies rotation, for +# s >= seq_len it copies through. This subsumes the `torch.cat([rotated, +# x[i, seq_len:]])` step (1 fewer kernel + 1 fewer alloc per call). Also +# skips the upstream `.contiguous()` because we no longer slice x. +@triton.jit +def _rope_apply_kernel( + x_ptr, # [S_total, N, D] bf16 (D is even, pairs are (a,b)=(2d, 2d+1)) + cos_ptr, # [seq_len, D/2] fp32 (valid only for s < seq_len) + sin_ptr, # [seq_len, D/2] fp32 + out_ptr, # [S_total, N, D] bf16 + SEQ_LEN, N, D_half, + x_stride_s, x_stride_n, + o_stride_s, o_stride_n, + cs_stride_s, + BLOCK_N: tl.constexpr, + BLOCK_D: tl.constexpr, +): + pid_s = tl.program_id(0) + pid_n = tl.program_id(1) + + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, BLOCK_D) # over the D/2 pairs + + n_mask = offs_n < N + d_mask = offs_d < D_half + + x_row_base = pid_s * x_stride_s + o_row_base = pid_s * o_stride_s + a_offs = x_row_base + offs_n[:, None] * x_stride_n + (2 * offs_d)[None, :] + b_offs = a_offs + 1 + mask = n_mask[:, None] & d_mask[None, :] + a = tl.load(x_ptr + a_offs, mask=mask, other=0.0).to(tl.float32) + b = tl.load(x_ptr + b_offs, mask=mask, other=0.0).to(tl.float32) + + a_out_offs = o_row_base + offs_n[:, None] * o_stride_n + (2 * offs_d)[None, :] + b_out_offs = a_out_offs + 1 + + if pid_s < SEQ_LEN: + cs_base = pid_s * cs_stride_s + cos = tl.load(cos_ptr + cs_base + offs_d, mask=d_mask, other=0.0).to(tl.float32) + sin = tl.load(sin_ptr + cs_base + offs_d, mask=d_mask, other=0.0).to(tl.float32) + # Rotate: (a + bi) * (cos + sin i) = (a*cos - b*sin) + (a*sin + b*cos) i + out_a = a * cos[None, :] - b * sin[None, :] + out_b = a * sin[None, :] + b * cos[None, :] + tl.store(out_ptr + a_out_offs, out_a, mask=mask) + tl.store(out_ptr + b_out_offs, out_b, mask=mask) + else: + # passthrough copy for the unrotated tail + tl.store(out_ptr + a_out_offs, a, mask=mask) + tl.store(out_ptr + b_out_offs, b, mask=mask) + + +def _split_complex_to_cos_sin(freqs_complex: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Convert complex128 freqs to (cos_f32, sin_f32) — once per cache entry.""" + # freqs_complex shape: (S, 1, D/2). Squeeze the middle 1. + if freqs_complex.dim() == 3 and freqs_complex.size(1) == 1: + freqs_complex = freqs_complex.squeeze(1) + cos = freqs_complex.real.to(torch.float32).contiguous() + sin = freqs_complex.imag.to(torch.float32).contiguous() + return cos, sin + + +def rope_apply_triton( + x: torch.Tensor, # [S_total, N, D] bf16 (or fp16/fp32) + cos_f32: torch.Tensor, # [seq_len, D/2] fp32 + sin_f32: torch.Tensor, # [seq_len, D/2] fp32 + seq_len: int | None = None, +) -> torch.Tensor: + """Apply rotary embedding via Triton kernel. + + iter-46: when `seq_len < x.size(0)`, the kernel rotates the first + `seq_len` rows and copies through rows `[seq_len:]`. This replaces the + `cat([rotated, x[i, seq_len:]])` pattern in `causal_rope_apply` with a + single kernel + single allocation. `seq_len=None` (default) means rotate + all rows (equivalent to iter-42 behavior). + + Returns a tensor of the same shape and dtype as `x`. + """ + assert x.dim() == 3, f"expected x.shape == (S, N, D), got {x.shape}" + S_total, N, D = x.shape + assert D % 2 == 0 + D_half = D // 2 + if seq_len is None: + seq_len = S_total + assert seq_len <= S_total + assert cos_f32.shape == (seq_len, D_half), \ + f"cos_f32 expected ({seq_len},{D_half}), got {cos_f32.shape}" + assert sin_f32.shape == (seq_len, D_half) + + out = torch.empty_like(x) + + BLOCK_N = 8 + BLOCK_D = triton.next_power_of_2(D_half) + grid = (S_total, triton.cdiv(N, BLOCK_N)) + + _rope_apply_kernel[grid]( + x, cos_f32, sin_f32, out, + seq_len, N, D_half, + x.stride(0), x.stride(1), + out.stride(0), out.stride(1), + cos_f32.stride(0), + BLOCK_N=BLOCK_N, BLOCK_D=BLOCK_D, + ) + return out diff --git a/utils/scheduler.py b/utils/scheduler.py new file mode 100644 index 0000000..cde3f85 --- /dev/null +++ b/utils/scheduler.py @@ -0,0 +1,194 @@ +from abc import abstractmethod, ABC +import torch + + +class SchedulerInterface(ABC): + """ + Base class for diffusion noise schedule. + """ + alphas_cumprod: torch.Tensor # [T], alphas for defining the noise schedule + + @abstractmethod + def add_noise( + self, clean_latent: torch.Tensor, + noise: torch.Tensor, timestep: torch.Tensor + ): + """ + Diffusion forward corruption process. + Input: + - clean_latent: the clean latent with shape [B, C, H, W] + - noise: the noise with shape [B, C, H, W] + - timestep: the timestep with shape [B] + Output: the corrupted latent with shape [B, C, H, W] + """ + pass + + def convert_x0_to_noise( + self, x0: torch.Tensor, xt: torch.Tensor, + timestep: torch.Tensor + ) -> torch.Tensor: + """ + Convert the diffusion network's x0 prediction to noise predidction. + x0: the predicted clean data with shape [B, C, H, W] + xt: the input noisy data with shape [B, C, H, W] + timestep: the timestep with shape [B] + + noise = (xt-sqrt(alpha_t)*x0) / sqrt(beta_t) (eq 11 in https://arxiv.org/abs/2311.18828) + """ + # use higher precision for calculations + original_dtype = x0.dtype + x0, xt, alphas_cumprod = map( + lambda x: x.double().to(x0.device), [x0, xt, + self.alphas_cumprod] + ) + + alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1) + beta_prod_t = 1 - alpha_prod_t + + noise_pred = (xt - alpha_prod_t ** + (0.5) * x0) / beta_prod_t ** (0.5) + return noise_pred.to(original_dtype) + + def convert_noise_to_x0( + self, noise: torch.Tensor, xt: torch.Tensor, + timestep: torch.Tensor + ) -> torch.Tensor: + """ + Convert the diffusion network's noise prediction to x0 predidction. + noise: the predicted noise with shape [B, C, H, W] + xt: the input noisy data with shape [B, C, H, W] + timestep: the timestep with shape [B] + + x0 = (x_t - sqrt(beta_t) * noise) / sqrt(alpha_t) (eq 11 in https://arxiv.org/abs/2311.18828) + """ + # use higher precision for calculations + original_dtype = noise.dtype + noise, xt, alphas_cumprod = map( + lambda x: x.double().to(noise.device), [noise, xt, + self.alphas_cumprod] + ) + alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1) + beta_prod_t = 1 - alpha_prod_t + + x0_pred = (xt - beta_prod_t ** + (0.5) * noise) / alpha_prod_t ** (0.5) + return x0_pred.to(original_dtype) + + def convert_velocity_to_x0( + self, velocity: torch.Tensor, xt: torch.Tensor, + timestep: torch.Tensor + ) -> torch.Tensor: + """ + Convert the diffusion network's velocity prediction to x0 predidction. + velocity: the predicted noise with shape [B, C, H, W] + xt: the input noisy data with shape [B, C, H, W] + timestep: the timestep with shape [B] + + v = sqrt(alpha_t) * noise - sqrt(beta_t) x0 + noise = (xt-sqrt(alpha_t)*x0) / sqrt(beta_t) + given v, x_t, we have + x0 = sqrt(alpha_t) * x_t - sqrt(beta_t) * v + see derivations https://chatgpt.com/share/679fb6c8-3a30-8008-9b0e-d1ae892dac56 + """ + # use higher precision for calculations + original_dtype = velocity.dtype + velocity, xt, alphas_cumprod = map( + lambda x: x.double().to(velocity.device), [velocity, xt, + self.alphas_cumprod] + ) + alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1) + beta_prod_t = 1 - alpha_prod_t + + x0_pred = (alpha_prod_t ** 0.5) * xt - (beta_prod_t ** 0.5) * velocity + return x0_pred.to(original_dtype) + + +class FlowMatchScheduler(): + + def __init__(self, num_inference_steps=100, num_train_timesteps=1000, shift=3.0, sigma_max=1.0, sigma_min=0.003 / 1.002, inverse_timesteps=False, extra_one_step=False, reverse_sigmas=False): + self.num_train_timesteps = num_train_timesteps + self.shift = shift + self.sigma_max = sigma_max + self.sigma_min = sigma_min + self.inverse_timesteps = inverse_timesteps + self.extra_one_step = extra_one_step + self.reverse_sigmas = reverse_sigmas + self.set_timesteps(num_inference_steps) + + def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, training=False): + sigma_start = self.sigma_min + \ + (self.sigma_max - self.sigma_min) * denoising_strength + if self.extra_one_step: + self.sigmas = torch.linspace( + sigma_start, self.sigma_min, num_inference_steps + 1)[:-1] + else: + self.sigmas = torch.linspace( + sigma_start, self.sigma_min, num_inference_steps) + if self.inverse_timesteps: + self.sigmas = torch.flip(self.sigmas, dims=[0]) + self.sigmas = self.shift * self.sigmas / \ + (1 + (self.shift - 1) * self.sigmas) + if self.reverse_sigmas: + self.sigmas = 1 - self.sigmas + self.timesteps = self.sigmas * self.num_train_timesteps + if training: + x = self.timesteps + y = torch.exp(-2 * ((x - num_inference_steps / 2) / + num_inference_steps) ** 2) + y_shifted = y - y.min() + bsmntw_weighing = y_shifted * \ + (num_inference_steps / y_shifted.sum()) + self.linear_timesteps_weights = bsmntw_weighing + + def step(self, model_output, timestep, sample, to_final=False): + if timestep.ndim == 2: + timestep = timestep.flatten(0, 1) + self.sigmas = self.sigmas.to(model_output.device) + self.timesteps = self.timesteps.to(model_output.device) + timestep_id = torch.argmin( + (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1) + if to_final or (timestep_id + 1 >= len(self.timesteps)).any(): + sigma_ = 1 if ( + self.inverse_timesteps or self.reverse_sigmas) else 0 + else: + sigma_ = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1) + prev_sample = sample + model_output * (sigma_ - sigma) + return prev_sample + + def add_noise(self, original_samples, noise, timestep): + """ + Diffusion forward corruption process. + Input: + - clean_latent: the clean latent with shape [B*T, C, H, W] + - noise: the noise with shape [B*T, C, H, W] + - timestep: the timestep with shape [B*T] + Output: the corrupted latent with shape [B*T, C, H, W] + """ + if timestep.ndim == 2: + timestep = timestep.flatten(0, 1) + self.sigmas = self.sigmas.to(noise.device) + self.timesteps = self.timesteps.to(noise.device) + timestep_id = torch.argmin( + (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1) + sample = (1 - sigma) * original_samples + sigma * noise + return sample.type_as(noise) + + def training_target(self, sample, noise, timestep): + target = noise - sample + return target + + def training_weight(self, timestep): + """ + Input: + - timestep: the timestep with shape [B*T] + Output: the corresponding weighting [B*T] + """ + if timestep.ndim == 2: + timestep = timestep.flatten(0, 1) + self.linear_timesteps_weights = self.linear_timesteps_weights.to(timestep.device) + timestep_id = torch.argmin( + (self.timesteps.unsqueeze(1) - timestep.unsqueeze(0)).abs(), dim=0) + weights = self.linear_timesteps_weights[timestep_id] + return weights diff --git a/utils/torch_compile_utils.py b/utils/torch_compile_utils.py new file mode 100644 index 0000000..c1e5edb --- /dev/null +++ b/utils/torch_compile_utils.py @@ -0,0 +1,117 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0 +# +# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied. +# +# SPDX-License-Identifier: Apache-2.0 +import torch +import torch.distributed as dist + + +def _is_main_process() -> bool: + return not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0 + + +def _log_once(message: str) -> None: + if _is_main_process(): + print(message) + + +class SafeCompiledCallable: + """Lazy torch.compile wrapper that falls back to eager on compile/runtime errors.""" + + def __init__( + self, + fn, + *, + name: str, + backend: str = "inductor", + mode: str | None = "max-autotune-no-cudagraphs", + fullgraph: bool = False, + dynamic: bool | None = False, + options: dict | None = None, + suppress_errors: bool = True, + ) -> None: + self.fn = fn + self.name = name + self.enabled = True + self.failed = False + self.failure_reason = None + + if suppress_errors: + try: + import torch._dynamo as torch_dynamo + + torch_dynamo.config.suppress_errors = True + except Exception as exc: + _log_once(f"[torch.compile] Could not enable suppress_errors: {exc}") + + compile_kwargs = { + "backend": backend, + "fullgraph": fullgraph, + "dynamic": dynamic, + } + if mode: + compile_kwargs["mode"] = mode + if options: + compile_kwargs["options"] = options + + _log_once( + "[torch.compile] Preparing " + f"{name}: backend={backend}, mode={mode}, " + f"fullgraph={fullgraph}, dynamic={dynamic}" + ) + self.compiled_fn = torch.compile(fn, **compile_kwargs) + + def __call__(self, *args, **kwargs): + if not self.enabled: + return self.fn(*args, **kwargs) + + try: + return self.compiled_fn(*args, **kwargs) + except Exception as exc: + self.enabled = False + self.failed = True + self.failure_reason = repr(exc) + _log_once( + f"[torch.compile][warn] {self.name} failed; " + f"falling back to eager. reason={exc}" + ) + return self.fn(*args, **kwargs) + + +def configure_module_call_torch_compile( + module, + *, + name: str, + backend: str = "inductor", + mode: str | None = "max-autotune-no-cudagraphs", + fullgraph: bool = False, + dynamic: bool | None = False, + options: dict | None = None, + suppress_errors: bool = True, +): + if not torch.cuda.is_available(): + _log_once(f"[torch.compile] Skipping {name}: CUDA is not available") + return None + + try: + return SafeCompiledCallable( + module, + name=name, + backend=backend, + mode=mode, + fullgraph=fullgraph, + dynamic=dynamic, + options=options, + suppress_errors=suppress_errors, + ) + except Exception as exc: + _log_once( + f"[torch.compile][warn] Could not prepare {name}; " + f"continuing in eager mode. reason={exc}" + ) + return None diff --git a/utils/wan_5b_wrapper.py b/utils/wan_5b_wrapper.py new file mode 100644 index 0000000..4884cb5 --- /dev/null +++ b/utils/wan_5b_wrapper.py @@ -0,0 +1,582 @@ +import types +from typing import List, Optional +import os +import torch +from torch import nn + +from utils.scheduler import SchedulerInterface, FlowMatchScheduler + +from wan_5b.modules.tokenizers import HuggingfaceTokenizer +from wan_5b.modules.model import WanModel +from wan_5b.modules.vae2_2 import _video_vae +from wan_5b.modules.t5 import umt5_xxl +from wan_5b.modules.causal_model import CausalWanModel + + +class WanTextEncoder(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + + self.text_encoder = umt5_xxl( + encoder_only=True, + return_tokenizer=False, + dtype=torch.float32, + device=torch.device('cpu') + ).eval().requires_grad_(False) + self.text_encoder.load_state_dict( + torch.load("wan_models/Wan2.2-TI2V-5B/models_t5_umt5-xxl-enc-bf16.pth", + map_location='cpu', weights_only=False) + ) + + # Move text encoder to GPU if available + if torch.cuda.is_available(): + self.text_encoder = self.text_encoder.cuda() + + self.tokenizer = HuggingfaceTokenizer( + name="wan_models/Wan2.2-TI2V-5B/google/umt5-xxl/", seq_len=512, clean='whitespace') + + @property + def device(self): + # Assume we are always on GPU + return torch.cuda.current_device() + + def forward(self, text_prompts: List[str]) -> dict: + ids, mask = self.tokenizer( + text_prompts, return_mask=True, add_special_tokens=True) + ids = ids.to(self.device) + mask = mask.to(self.device) + seq_lens = mask.gt(0).sum(dim=1).long() + context = self.text_encoder(ids, mask) + for u, v in zip(context, seq_lens): + u[v:] = 0.0 # set padding to 0.0 + + return { + "prompt_embeds": context + } + + +class WanVAEWrapper(torch.nn.Module): + def __init__(self): + super().__init__() + mean = [ + -0.2289, + -0.0052, + -0.1323, + -0.2339, + -0.2799, + 0.0174, + 0.1838, + 0.1557, + -0.1382, + 0.0542, + 0.2813, + 0.0891, + 0.1570, + -0.0098, + 0.0375, + -0.1825, + -0.2246, + -0.1207, + -0.0698, + 0.5109, + 0.2665, + -0.2108, + -0.2158, + 0.2502, + -0.2055, + -0.0322, + 0.1109, + 0.1567, + -0.0729, + 0.0899, + -0.2799, + -0.1230, + -0.0313, + -0.1649, + 0.0117, + 0.0723, + -0.2839, + -0.2083, + -0.0520, + 0.3748, + 0.0152, + 0.1957, + 0.1433, + -0.2944, + 0.3573, + -0.0548, + -0.1681, + -0.0667, + ] + std = [ + 0.4765, + 1.0364, + 0.4514, + 1.1677, + 0.5313, + 0.4990, + 0.4818, + 0.5013, + 0.8158, + 1.0344, + 0.5894, + 1.0901, + 0.6885, + 0.6165, + 0.8454, + 0.4978, + 0.5759, + 0.3523, + 0.7135, + 0.6804, + 0.5833, + 1.4146, + 0.8986, + 0.5659, + 0.7069, + 0.5338, + 0.4889, + 0.4917, + 0.4069, + 0.4999, + 0.6866, + 0.4093, + 0.5709, + 0.6065, + 0.6415, + 0.4944, + 0.5726, + 1.2042, + 0.5458, + 1.6887, + 0.3971, + 1.0600, + 0.3943, + 0.5537, + 0.5444, + 0.4089, + 0.7468, + 0.7744, + ] + self.mean = torch.tensor(mean, dtype=torch.float32) + self.std = torch.tensor(std, dtype=torch.float32) + + # init model + self.model = _video_vae( + pretrained_path="wan_models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth", + ).eval().requires_grad_(False) + + def encode_to_latent(self, pixel: torch.Tensor) -> torch.Tensor: + # pixel: [batch_size, num_channels, num_frames, height, width] + device, dtype = pixel.device, pixel.dtype + + scale = [self.mean.to(device=device, dtype=dtype), + 1.0 / self.std.to(device=device, dtype=dtype)] + + output = [ + self.model.encode(u.unsqueeze(0), scale).float().squeeze(0) + for u in pixel + ] + output = torch.stack(output, dim=0) + # from [batch_size, num_channels, num_frames, height, width] + # to [batch_size, num_frames, num_channels, height, width] + output = output.permute(0, 2, 1, 3, 4) + return output + + def decode_to_pixel(self, latent: torch.Tensor, use_cache: bool = False) -> torch.Tensor: + # from [batch_size, num_frames, num_channels, height, width] + # to [batch_size, num_channels, num_frames, height, width] + zs = latent.permute(0, 2, 1, 3, 4) + if use_cache: + assert latent.shape[0] == 1, "Batch size must be 1 when using cache" + + device, dtype = latent.device, latent.dtype + scale = [self.mean.to(device=device, dtype=dtype), + 1.0 / self.std.to(device=device, dtype=dtype)] + + if use_cache: + decode_function = self.model.cached_decode + else: + decode_function = self.model.decode + + output = [] + for u in zs: + output.append(decode_function(u.unsqueeze(0), scale).float().clamp_(-1, 1).squeeze(0)) + output = torch.stack(output, dim=0) + # from [batch_size, num_channels, num_frames, height, width] + # to [batch_size, num_frames, num_channels, height, width] + output = output.permute(0, 2, 1, 3, 4) + return output + + def decode_to_pixel_chunk(self, latent: torch.Tensor, use_cache: bool = False, chunk_size: int = 1) -> torch.Tensor: + """ + Decode latent frames to pixel space. + + Args: + latent: Latent tensor with shape [batch_size, num_frames, num_channels, height, width] + use_cache: Whether to use cached decoding (for streaming) + chunk_size: Number of latent frames to decode at once (default 240 to avoid OOM) + + Returns: + Decoded video tensor with shape [batch_size, num_frames, num_channels, height, width] + """ + # latent shape: [batch_size, num_frames, num_channels, height, width] + # zs shape after permute: [batch_size, num_channels, num_frames, height, width] + zs = latent.permute(0, 2, 1, 3, 4) + if use_cache: + assert latent.shape[0] == 1, "Batch size must be 1 when using cache" + + device, dtype = latent.device, latent.dtype + scale = [self.mean.to(device=device, dtype=dtype), + 1.0 / self.std.to(device=device, dtype=dtype)] + + if use_cache: + decode_function = self.model.cached_decode + else: + decode_function = self.model.decode + + output = [] + for u in zs: + num_frames = u.shape[1] + if num_frames <= chunk_size: + # Decode short clips in one pass. + if use_cache: + # Start this segment from a clean cache. + self.model.clear_cache() + decoded = decode_function(u.unsqueeze(0), scale).float().clamp_(-1, 1).squeeze(0) + decoded = decoded.cpu() + if use_cache: + # Clear after this segment so it cannot affect the next video. + self.model.clear_cache() + else: + # Decode longer clips in temporal chunks. + decoded_chunks = [] + if use_cache: + # Clear once at the segment start; later chunks share the + # internal cache. + self.model.clear_cache() + for start_idx in range(0, num_frames, chunk_size): + end_idx = min(start_idx + chunk_size, num_frames) + chunk = u[:, start_idx:end_idx, :, :] # [C, chunk_frames, H, W] + decoded_chunk = decode_function(chunk.unsqueeze(0), scale).float().clamp_(-1, 1).squeeze(0) + decoded_chunks.append(decoded_chunk.cpu()) + + del decoded_chunk + torch.cuda.empty_cache() + decoded = torch.cat(decoded_chunks, dim=1) + if use_cache: + # Clear the cache after the full segment. + self.model.clear_cache() + output.append(decoded) + + output = torch.stack(output, dim=0) + output = output.permute(0, 2, 1, 3, 4) + return output + + +class WanDiffusionWrapper(torch.nn.Module): + def __init__( + self, + model_name="Wan2.2-TI2V-5B", + timestep_shift=8.0, + is_causal=False, + local_attn_size=-1, + sink_size=0, + num_frame_per_block=1, + t_scale=1.0, + rope_method="linear", + original_seq_len=None, + ): + super().__init__() + + if is_causal: + self.model = CausalWanModel.from_pretrained( + f"wan_models/{model_name}/", local_attn_size=local_attn_size, sink_size=sink_size, + num_frame_per_block=num_frame_per_block) + else: + self.model = WanModel.from_pretrained(f"wan_models/{model_name}/") + self.model.eval() + self.model.t_scale = t_scale + self.model.rope_method = rope_method + self.model.original_seq_len = original_seq_len + + # For non-causal diffusion, all frames share the same timestep + self.uniform_timestep = not is_causal + + self.scheduler = FlowMatchScheduler( + shift=timestep_shift, sigma_min=0.0, extra_one_step=True + ) + self.scheduler.set_timesteps(1000, training=True) + + self.seq_len = 28160 # [1, 32, 48, 44, 80] + + self.post_init() + self._compiled_model_call = None + + def enable_gradient_checkpointing(self) -> None: + self.model.enable_gradient_checkpointing() + + def configure_torch_compile( + self, + *, + backend: str = "inductor", + mode: str | None = "max-autotune-no-cudagraphs", + fullgraph: bool = False, + dynamic: bool | None = False, + options: dict | None = None, + suppress_errors: bool = True, + ) -> bool: + from utils.torch_compile_utils import configure_module_call_torch_compile + + self._compiled_model_call = configure_module_call_torch_compile( + self.model, + name="WanDiffusionWrapper5B.model", + backend=backend, + mode=mode, + fullgraph=fullgraph, + dynamic=dynamic, + options=options, + suppress_errors=suppress_errors, + ) + return self._compiled_model_call is not None + + def _call_model(self, *args, **kwargs): + # iter-39 v2: publish kv_cache scalars BEFORE entering the compiled + # graph. The earlier version (iter-39 v1) published them inside + # `_forward_inference`, but that function IS compiled, so each + # `.item()` triggered a graph break. Moving the reads to this eager + # wrapper keeps the dict lookups in the compiled attention forward + # free of `.item()` syncs without adding any graph break. + kv_cache = kwargs.get("kv_cache", None) + if kv_cache is not None and len(kv_cache) > 0: + try: + from wan_5b.modules.causal_model import _CURRENT_GRID_META + first_block_cache = kv_cache[0] + _CURRENT_GRID_META["global_end_index"] = int( + first_block_cache["global_end_index"].item() + ) + _CURRENT_GRID_META["local_end_index"] = int( + first_block_cache["local_end_index"].item() + ) + _ps = first_block_cache.get("pinned_start", None) + if _ps is not None and hasattr(_ps, "item"): + _CURRENT_GRID_META["pinned_start"] = int(_ps.item()) + _CURRENT_GRID_META["pinned_len"] = int( + first_block_cache["pinned_len"].item() + ) + else: + _CURRENT_GRID_META["pinned_start"] = -1 + _CURRENT_GRID_META["pinned_len"] = 0 + except (KeyError, AttributeError, ImportError): + pass + defer_kv_updates = ( + os.environ.get("LLV2_DEFER_KV_UPDATES", "0") == "1" + and kv_cache is not None + ) + if defer_kv_updates: + kwargs["defer_cache_updates"] = True + + if self._compiled_model_call is not None: + # iter-25: signal cudagraph allocator that a new "step" starts. + # Required for mode=reduce-overhead when modules cache state + # (KV cache rolling buffers, fp4-quant scale tensors) so the + # cudagraph pool knows it can safely reuse step-N memory now + # that step-(N+1) is starting. + mark_step = getattr(torch.compiler, "cudagraph_mark_step_begin", None) + if mark_step is not None: + mark_step() + result = self._compiled_model_call(*args, **kwargs) + else: + result = self.model(*args, **kwargs) + + if defer_kv_updates: + if not isinstance(result, tuple) or len(result) != 2: + raise RuntimeError( + "LLV2_DEFER_KV_UPDATES expected model to return " + "(output, cache_update_infos)." + ) + output, cache_update_infos = result + if cache_update_infos: + self.model._apply_cache_updates(kv_cache, cache_update_infos) + return output + return result + + def _convert_flow_pred_to_x0(self, flow_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: + """ + Convert flow matching's prediction to x0 prediction. + flow_pred: the prediction with shape [B, C, H, W] + xt: the input noisy data with shape [B, C, H, W] + timestep: the timestep with shape [B] + + pred = noise - x0 + x_t = (1-sigma_t) * x0 + sigma_t * noise + we have x0 = x_t - sigma_t * pred + see derivations https://chatgpt.com/share/67bf8589-3d04-8008-bc6e-4cf1a24e2d0e + """ + # use higher precision for calculations + original_dtype = flow_pred.dtype + flow_pred, xt, sigmas, timesteps = map( + lambda x: x.double().to(flow_pred.device), [flow_pred, xt, + self.scheduler.sigmas, + self.scheduler.timesteps] + ) + + timestep_id = torch.argmin( + (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1) + x0_pred = xt - sigma_t * flow_pred + return x0_pred.to(original_dtype) + + @staticmethod + def _convert_x0_to_flow_pred(scheduler, x0_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: + """ + Convert x0 prediction to flow matching's prediction. + x0_pred: the x0 prediction with shape [B, C, H, W] + xt: the input noisy data with shape [B, C, H, W] + timestep: the timestep with shape [B] + + pred = (x_t - x_0) / sigma_t + """ + # use higher precision for calculations + original_dtype = x0_pred.dtype + x0_pred, xt, sigmas, timesteps = map( + lambda x: x.double().to(x0_pred.device), [x0_pred, xt, + scheduler.sigmas, + scheduler.timesteps] + ) + timestep_id = torch.argmin( + (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1) + flow_pred = (xt - x0_pred) / sigma_t + return flow_pred.to(original_dtype) + + def forward( + self, + noisy_image_or_video: torch.Tensor, conditional_dict: dict, + timestep: torch.Tensor, kv_cache: Optional[List[dict]] = None, + crossattn_cache: Optional[List[dict]] = None, + current_start: Optional[int] = None, + classify_mode: Optional[bool] = False, + concat_time_embeddings: Optional[bool] = False, + clean_x: Optional[torch.Tensor] = None, + aug_t: Optional[torch.Tensor] = None, + cache_start: Optional[int] = None, + rope_temporal_offset: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + prompt_embeds = conditional_dict["prompt_embeds"] + + # [B, F] -> [B] + if self.uniform_timestep: + input_timestep = timestep[:, 0] + else: + input_timestep = timestep + + logits = None + rope_offset_was_set = ( + rope_temporal_offset is not None + and hasattr(self.model, "rope_temporal_offset") + ) + if rope_offset_was_set: + prev_rope_temporal_offset = self.model.rope_temporal_offset + self.model.rope_temporal_offset = rope_temporal_offset + + # X0 prediction + if kv_cache is not None: + flow_pred = self._call_model( + noisy_image_or_video.permute(0, 2, 1, 3, 4), + t=input_timestep, context=prompt_embeds, + seq_len=self.seq_len, + kv_cache=kv_cache, + crossattn_cache=crossattn_cache, + current_start=current_start, + cache_start=cache_start + ).permute(0, 2, 1, 3, 4) + else: + if clean_x is not None: + # teacher forcing + flow_pred = self._call_model( + noisy_image_or_video.permute(0, 2, 1, 3, 4), + t=input_timestep, context=prompt_embeds, + seq_len=self.seq_len, + clean_x=clean_x.permute(0, 2, 1, 3, 4), + aug_t=aug_t, + ).permute(0, 2, 1, 3, 4) + else: + if classify_mode: + flow_pred, logits = self._call_model( + noisy_image_or_video.permute(0, 2, 1, 3, 4), + t=input_timestep, context=prompt_embeds, + seq_len=self.seq_len, + classify_mode=True, + register_tokens=self._register_tokens, + cls_pred_branch=self._cls_pred_branch, + gan_ca_blocks=self._gan_ca_blocks, + concat_time_embeddings=concat_time_embeddings + ) + flow_pred = flow_pred.permute(0, 2, 1, 3, 4) + else: + flow_pred = self._call_model( + noisy_image_or_video.permute(0, 2, 1, 3, 4), + t=input_timestep, context=prompt_embeds, + seq_len=self.seq_len + ).permute(0, 2, 1, 3, 4) + + if rope_offset_was_set: + self.model.rope_temporal_offset = prev_rope_temporal_offset + + pred_x0 = self._convert_flow_pred_to_x0( + flow_pred=flow_pred.flatten(0, 1), + xt=noisy_image_or_video.flatten(0, 1), + timestep=timestep.flatten(0, 1) + ).unflatten(0, flow_pred.shape[:2]) + + if logits is not None: + return flow_pred, pred_x0, logits + + return flow_pred, pred_x0 + + def get_scheduler(self) -> SchedulerInterface: + """ + Update the current scheduler with the interface's static method + """ + scheduler = self.scheduler + scheduler.convert_x0_to_noise = types.MethodType( + SchedulerInterface.convert_x0_to_noise, scheduler) + scheduler.convert_noise_to_x0 = types.MethodType( + SchedulerInterface.convert_noise_to_x0, scheduler) + scheduler.convert_velocity_to_x0 = types.MethodType( + SchedulerInterface.convert_velocity_to_x0, scheduler) + self.scheduler = scheduler + return scheduler + + def post_init(self): + """ + A few custom initialization steps that should be called after the object is created. + Currently, the only one we have is to bind a few methods to scheduler. + We can gradually add more methods here if needed. + """ + self.get_scheduler() + + +_MG_LIGHTVAE_DEFAULT_PATHS = { + "mg_lightvae": os.path.join("wan_models", "Matrix-Game-3.0", "MG-LightVAE.pth"), + "mg_lightvae_v2": os.path.join("wan_models", "Matrix-Game-3.0", "MG-LightVAE_v2.pth"), +} + + +def build_vae_5b(args): + """Return the 5B VAE wrapper requested by args.vae_type.""" + vae_type = str(getattr(args, "vae_type", "wan")).lower().strip() + + if vae_type in ("wan", "wan2.2", ""): + return WanVAEWrapper() + + if vae_type in _MG_LIGHTVAE_DEFAULT_PATHS: + from utils.lightvae_5b_wrapper import LightVAE5BWrapper + + return LightVAE5BWrapper(vae_path=_MG_LIGHTVAE_DEFAULT_PATHS[vae_type]) + + raise ValueError( + f"Unknown vae_type '{vae_type}'. " + "Expected one of: wan, mg_lightvae, mg_lightvae_v2." + ) diff --git a/wan_5b/__init__.py b/wan_5b/__init__.py new file mode 100644 index 0000000..ac59c4f --- /dev/null +++ b/wan_5b/__init__.py @@ -0,0 +1,16 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. + +__all__ = ["WanI2V", "WanT2V", "WanTI2V"] + + +def __getattr__(name): + if name == "WanI2V": + from .image2video import WanI2V + return WanI2V + if name == "WanT2V": + from .text2video import WanT2V + return WanT2V + if name == "WanTI2V": + from .textimage2video import WanTI2V + return WanTI2V + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/wan_5b/configs/__init__.py b/wan_5b/configs/__init__.py new file mode 100644 index 0000000..875afe7 --- /dev/null +++ b/wan_5b/configs/__init__.py @@ -0,0 +1,39 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import copy +import os + +os.environ['TOKENIZERS_PARALLELISM'] = 'false' + +from .wan_i2v_A14B import i2v_A14B +from .wan_t2v_A14B import t2v_A14B +from .wan_ti2v_5B import ti2v_5B + +WAN_CONFIGS = { + 't2v-A14B': t2v_A14B, + 'i2v-A14B': i2v_A14B, + 'ti2v-5B': ti2v_5B, +} + +SIZE_CONFIGS = { + '720*1280': (720, 1280), + '1280*720': (1280, 720), + '480*832': (480, 832), + '832*480': (832, 480), + '704*1280': (704, 1280), + '1280*704': (1280, 704) +} + +MAX_AREA_CONFIGS = { + '720*1280': 720 * 1280, + '1280*720': 1280 * 720, + '480*832': 480 * 832, + '832*480': 832 * 480, + '704*1280': 704 * 1280, + '1280*704': 1280 * 704, +} + +SUPPORTED_SIZES = { + 't2v-A14B': ('720*1280', '1280*720', '480*832', '832*480'), + 'i2v-A14B': ('720*1280', '1280*720', '480*832', '832*480'), + 'ti2v-5B': ('704*1280', '1280*704'), +} diff --git a/wan_5b/configs/shared_config.py b/wan_5b/configs/shared_config.py new file mode 100644 index 0000000..c58ab04 --- /dev/null +++ b/wan_5b/configs/shared_config.py @@ -0,0 +1,20 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import torch +from easydict import EasyDict + +#------------------------ Wan shared config ------------------------# +wan_shared_cfg = EasyDict() + +# t5 +wan_shared_cfg.t5_model = 'umt5_xxl' +wan_shared_cfg.t5_dtype = torch.bfloat16 +wan_shared_cfg.text_len = 512 + +# transformer +wan_shared_cfg.param_dtype = torch.bfloat16 + +# inference +wan_shared_cfg.num_train_timesteps = 1000 +wan_shared_cfg.sample_fps = 16 +wan_shared_cfg.sample_neg_prompt = '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走' +wan_shared_cfg.frame_num = 81 diff --git a/wan_5b/configs/wan_i2v_A14B.py b/wan_5b/configs/wan_i2v_A14B.py new file mode 100644 index 0000000..f654cc6 --- /dev/null +++ b/wan_5b/configs/wan_i2v_A14B.py @@ -0,0 +1,37 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import torch +from easydict import EasyDict + +from .shared_config import wan_shared_cfg + +#------------------------ Wan I2V A14B ------------------------# + +i2v_A14B = EasyDict(__name__='Config: Wan I2V A14B') +i2v_A14B.update(wan_shared_cfg) + +i2v_A14B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth' +i2v_A14B.t5_tokenizer = 'google/umt5-xxl' + +# vae +i2v_A14B.vae_checkpoint = 'Wan2.1_VAE.pth' +i2v_A14B.vae_stride = (4, 8, 8) + +# transformer +i2v_A14B.patch_size = (1, 2, 2) +i2v_A14B.dim = 5120 +i2v_A14B.ffn_dim = 13824 +i2v_A14B.freq_dim = 256 +i2v_A14B.num_heads = 40 +i2v_A14B.num_layers = 40 +i2v_A14B.window_size = (-1, -1) +i2v_A14B.qk_norm = True +i2v_A14B.cross_attn_norm = True +i2v_A14B.eps = 1e-6 +i2v_A14B.low_noise_checkpoint = 'low_noise_model' +i2v_A14B.high_noise_checkpoint = 'high_noise_model' + +# inference +i2v_A14B.sample_shift = 5.0 +i2v_A14B.sample_steps = 40 +i2v_A14B.boundary = 0.900 +i2v_A14B.sample_guide_scale = (3.5, 3.5) # low noise, high noise diff --git a/wan_5b/configs/wan_t2v_A14B.py b/wan_5b/configs/wan_t2v_A14B.py new file mode 100644 index 0000000..a5220a5 --- /dev/null +++ b/wan_5b/configs/wan_t2v_A14B.py @@ -0,0 +1,37 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +from easydict import EasyDict + +from .shared_config import wan_shared_cfg + +#------------------------ Wan T2V A14B ------------------------# + +t2v_A14B = EasyDict(__name__='Config: Wan T2V A14B') +t2v_A14B.update(wan_shared_cfg) + +# t5 +t2v_A14B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth' +t2v_A14B.t5_tokenizer = 'google/umt5-xxl' + +# vae +t2v_A14B.vae_checkpoint = 'Wan2.1_VAE.pth' +t2v_A14B.vae_stride = (4, 8, 8) + +# transformer +t2v_A14B.patch_size = (1, 2, 2) +t2v_A14B.dim = 5120 +t2v_A14B.ffn_dim = 13824 +t2v_A14B.freq_dim = 256 +t2v_A14B.num_heads = 40 +t2v_A14B.num_layers = 40 +t2v_A14B.window_size = (-1, -1) +t2v_A14B.qk_norm = True +t2v_A14B.cross_attn_norm = True +t2v_A14B.eps = 1e-6 +t2v_A14B.low_noise_checkpoint = 'low_noise_model' +t2v_A14B.high_noise_checkpoint = 'high_noise_model' + +# inference +t2v_A14B.sample_shift = 12.0 +t2v_A14B.sample_steps = 40 +t2v_A14B.boundary = 0.875 +t2v_A14B.sample_guide_scale = (3.0, 4.0) # low noise, high noise diff --git a/wan_5b/configs/wan_ti2v_5B.py b/wan_5b/configs/wan_ti2v_5B.py new file mode 100644 index 0000000..d5d5aed --- /dev/null +++ b/wan_5b/configs/wan_ti2v_5B.py @@ -0,0 +1,36 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +from easydict import EasyDict + +from .shared_config import wan_shared_cfg + +#------------------------ Wan TI2V 5B ------------------------# + +ti2v_5B = EasyDict(__name__='Config: Wan TI2V 5B') +ti2v_5B.update(wan_shared_cfg) + +# t5 +ti2v_5B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth' +ti2v_5B.t5_tokenizer = 'google/umt5-xxl' + +# vae +ti2v_5B.vae_checkpoint = 'Wan2.2_VAE.pth' +ti2v_5B.vae_stride = (4, 16, 16) + +# transformer +ti2v_5B.patch_size = (1, 2, 2) +ti2v_5B.dim = 3072 +ti2v_5B.ffn_dim = 14336 +ti2v_5B.freq_dim = 256 +ti2v_5B.num_heads = 24 +ti2v_5B.num_layers = 30 +ti2v_5B.window_size = (-1, -1) +ti2v_5B.qk_norm = True +ti2v_5B.cross_attn_norm = True +ti2v_5B.eps = 1e-6 + +# inference +ti2v_5B.sample_fps = 24 +ti2v_5B.sample_shift = 5.0 +ti2v_5B.sample_steps = 50 +ti2v_5B.sample_guide_scale = 5.0 +ti2v_5B.frame_num = 121 diff --git a/wan_5b/distributed/__init__.py b/wan_5b/distributed/__init__.py new file mode 100644 index 0000000..566f71e --- /dev/null +++ b/wan_5b/distributed/__init__.py @@ -0,0 +1 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. diff --git a/wan_5b/distributed/fsdp.py b/wan_5b/distributed/fsdp.py new file mode 100644 index 0000000..6bb496d --- /dev/null +++ b/wan_5b/distributed/fsdp.py @@ -0,0 +1,43 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import gc +from functools import partial + +import torch +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +from torch.distributed.fsdp import MixedPrecision, ShardingStrategy +from torch.distributed.fsdp.wrap import lambda_auto_wrap_policy +from torch.distributed.utils import _free_storage + + +def shard_model( + model, + device_id, + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + buffer_dtype=torch.float32, + process_group=None, + sharding_strategy=ShardingStrategy.FULL_SHARD, + sync_module_states=True, +): + model = FSDP( + module=model, + process_group=process_group, + sharding_strategy=sharding_strategy, + auto_wrap_policy=partial( + lambda_auto_wrap_policy, lambda_fn=lambda m: m in model.blocks), + mixed_precision=MixedPrecision( + param_dtype=param_dtype, + reduce_dtype=reduce_dtype, + buffer_dtype=buffer_dtype), + device_id=device_id, + sync_module_states=sync_module_states) + return model + + +def free_model(model): + for m in model.modules(): + if isinstance(m, FSDP): + _free_storage(m._handle.flat_param.data) + del model + gc.collect() + torch.cuda.empty_cache() diff --git a/wan_5b/distributed/sequence_parallel.py b/wan_5b/distributed/sequence_parallel.py new file mode 100644 index 0000000..29118e4 --- /dev/null +++ b/wan_5b/distributed/sequence_parallel.py @@ -0,0 +1,477 @@ +# Adopted from https://github.com/Wan-Video/Wan2.2 +# SPDX-License-Identifier: Apache-2.0 +import torch +import torch.cuda.amp as amp + +from ..modules.model import sinusoidal_embedding_1d +from .ulysses import distributed_attention, distributed_flex_attention +from .util import gather_forward, get_rank, get_world_size +import math + + +def pad_freqs(original_tensor, target_len): + seq_len, s1, s2 = original_tensor.shape + pad_size = target_len - seq_len + padding_tensor = torch.ones( + pad_size, + s1, + s2, + dtype=original_tensor.dtype, + device=original_tensor.device) + padded_tensor = torch.cat([original_tensor, padding_tensor], dim=0) + return padded_tensor + + +from utils.position_embedding_utils import ( + compute_temporal_freqs as _compute_temporal_freqs, + select_temporal_offset_for_sample, +) + + +@torch.amp.autocast('cuda', enabled=False) +def sp_rope_apply( + x, + grid_sizes, + freqs, + t_scale=1.0, + method="linear", + original_seq_len=None, + temporal_offset=0.0, +): + """ + x: [B, L, N, C]. + grid_sizes: [B, 3]. + freqs: [M, C // 2]. + """ + n, c = x.size(2), x.size(3) // 2 + # split freqs + freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) + + # loop over samples + output = [] + for i, (f, h, w) in enumerate(grid_sizes.tolist()): + local_f = f + sp_rank = get_rank() + start_frame = sp_rank * local_f + seq_len = local_f * h * w + + # precompute multipliers + x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape( + seq_len, n, -1, 2)) + temporal_offset_i = select_temporal_offset_for_sample( + temporal_offset, i, local_f, start_frame=start_frame) + temporal_freqs = _compute_temporal_freqs( + freqs[0], local_f, start_frame, t_scale, x.device, + method=method, original_seq_len=original_seq_len, + temporal_offset=temporal_offset_i) + freqs_i = torch.cat([ + temporal_freqs.view(local_f, 1, 1, -1).expand(local_f, h, w, -1), + freqs[1][:h].view(1, h, 1, -1).expand(local_f, h, w, -1), + freqs[2][:w].view(1, 1, w, -1).expand(local_f, h, w, -1) + ], + dim=-1).reshape(seq_len, 1, -1) + + # apply rotary embedding + x_i = torch.view_as_real(x_i * freqs_i).flatten(2) + x_i = torch.cat([x_i, x[i, seq_len:]]) + + # append to collection + output.append(x_i) + return torch.stack(output).float() + + +def sp_dit_forward( + self, + x, + t, + context, + seq_len, + y=None, +): + """ + x: A list of videos each with shape [C, T, H, W]. + t: [B]. + context: A list of text embeddings each with shape [L, C]. + """ + if self.model_type == 'i2v': + assert y is not None + # params + device = self.patch_embedding.weight.device + if self.freqs.device != device: + self.freqs = self.freqs.to(device) + + if y is not None: + x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)] + + # embeddings + x = [self.patch_embedding(u.unsqueeze(0)) for u in x] + grid_sizes = torch.stack( + [torch.tensor(u.shape[2:], dtype=torch.long) for u in x]) + x = [u.flatten(2).transpose(1, 2) for u in x] + seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long) + assert seq_lens.max() <= seq_len + x = torch.cat([ + torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1) + for u in x + ]) + + # time embeddings + if t.dim() == 1: + t = t.expand(t.size(0), seq_len) + with torch.amp.autocast('cuda', dtype=torch.float32): + bt = t.size(0) + t = t.flatten() + e = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, + t).unflatten(0, (bt, seq_len)).float()) + e0 = self.time_projection(e).unflatten(2, (6, self.dim)) + assert e.dtype == torch.float32 and e0.dtype == torch.float32 + + # context + context_lens = None + context = self.text_embedding( + torch.stack([ + torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) + for u in context + ])) + + # Context Parallel + x = torch.chunk(x, get_world_size(), dim=1)[get_rank()] + e = torch.chunk(e, get_world_size(), dim=1)[get_rank()] + e0 = torch.chunk(e0, get_world_size(), dim=1)[get_rank()] + + # arguments + kwargs = dict( + e=e0, + seq_lens=seq_lens, + grid_sizes=grid_sizes, + freqs=self.freqs, + context=context, + context_lens=context_lens) + + for block in self.blocks: + x = block(x, **kwargs) + + # head + x = self.head(x, e) + + # Context Parallel + x = gather_forward(x, dim=1) + + # unpatchify + x = self.unpatchify(x, grid_sizes) + return [u.float() for u in x] + + +def sp_attn_forward(self, x, seq_lens, grid_sizes, freqs, dtype=torch.bfloat16, + t_scale=1.0, method="linear", original_seq_len=None, + temporal_offset=0.0): + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + half_dtypes = (torch.float16, torch.bfloat16) + + def half(x): + return x if x.dtype in half_dtypes else x.to(dtype) + + # query, key, value function + def qkv_fn(x): + q = self.norm_q(self.q(x)).view(b, s, n, d) + k = self.norm_k(self.k(x)).view(b, s, n, d) + v = self.v(x).view(b, s, n, d) + return q, k, v + + q, k, v = qkv_fn(x) + q = sp_rope_apply(q, grid_sizes, freqs, t_scale=t_scale, + method=method, original_seq_len=original_seq_len, + temporal_offset=temporal_offset) + k = sp_rope_apply(k, grid_sizes, freqs, t_scale=t_scale, + method=method, original_seq_len=original_seq_len, + temporal_offset=temporal_offset) + + x = distributed_attention( + half(q), + half(k), + half(v), + seq_lens, + window_size=self.window_size, + ) + + # output + x = x.flatten(2) + x = self.o(x) + return x + + + + +def sp_dit_causal_forward_train( + self, + x, + t, + context, + seq_len, + clean_x=None, + aug_t=None, + clip_fea=None, + y=None, +): + r""" + Forward pass through the diffusion model + + Args: + x (List[Tensor]): + List of input video tensors, each with shape [C_in, F, H, W] + t (Tensor): + Diffusion timesteps tensor of shape [B] + context (List[Tensor]): + List of text embeddings each with shape [L, C] + seq_len (`int`): + Maximum sequence length for positional encoding + clip_fea (Tensor, *optional*): + CLIP image features for image-to-video mode + y (List[Tensor], *optional*): + Conditional video inputs for image-to-video mode, same shape as x + + Returns: + List[Tensor]: + List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8] + """ + if self.model_type == 'i2v': + assert clip_fea is not None and y is not None + # params + device = self.patch_embedding.weight.device + if self.freqs.device != device: + self.freqs = self.freqs.to(device) + + # Construct the blockwise causal attention mask. Frames are sharded across + # SP ranks, so total frames = local frames per rank * sp_size. + sp_size = get_world_size() + # Recreate mask when batch size changes to avoid Triton broadcasting bug + current_batch_size = x.shape[0] + if self.block_mask is None or self._block_mask_batch_size != current_batch_size: + self._block_mask_batch_size = current_batch_size + if clean_x is not None: + if self.independent_first_frame: + raise NotImplementedError() + else: + self.block_mask = self._prepare_teacher_forcing_mask_natural( + device, num_frames=x.shape[2] * sp_size, + frame_seqlen=x.shape[-2] * x.shape[-1] // (self.patch_size[1] * self.patch_size[2]), + num_frame_per_block=self.num_frame_per_block, + sp_size=sp_size, + batch_size=current_batch_size, + ) + else: + if self.independent_first_frame: + self.block_mask = self._prepare_blockwise_causal_attn_mask_i2v( + device, num_frames=x.shape[2] * sp_size, + frame_seqlen=x.shape[-2] * x.shape[-1] // (self.patch_size[1] * self.patch_size[2]), + num_frame_per_block=self.num_frame_per_block, + batch_size=current_batch_size, + ) + else: + self.block_mask = self._prepare_blockwise_causal_attn_mask( + device, num_frames=x.shape[2] * sp_size, + frame_seqlen=x.shape[-2] * x.shape[-1] // (self.patch_size[1] * self.patch_size[2]), + num_frame_per_block=self.num_frame_per_block, + batch_size=current_batch_size, + ) + + if y is not None: + x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)] + + # embeddings + x = [self.patch_embedding(u.unsqueeze(0)) for u in x] + + grid_sizes = torch.stack( + [torch.tensor(u.shape[2:], dtype=torch.long) for u in x]) + x = [u.flatten(2).transpose(1, 2) for u in x] + + seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long) + max_len = int(seq_lens.max().item()) + assert max_len > 0, "Token sequence length is zero after patch embedding" + # Pad all samples to the batch max length instead of the first sample length + x = torch.cat([ + torch.cat([u, u.new_zeros(1, max_len - u.size(1), u.size(2))], dim=1) + for u in x + ]) + + # time embeddings + if t.dim() == 1: + raise NotImplementedError(f"t.shape should be [B, F], but got {t.shape}") + bt = t.size(0) + t_len = t.size(1) + t_ori_shape = t.shape + t = t.flatten() + e = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, t).unflatten(0, (bt, t_len)).type_as(x)) + e0 = self.time_projection(e).unflatten(2, (6, self.dim)) # B, F, 6, C + + # context + context_lens = None + context = self.text_embedding( + torch.stack([ + torch.cat( + [u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) + for u in context + ])) + + if clean_x is not None: + clean_x = [self.patch_embedding(u.unsqueeze(0)) for u in clean_x] + clean_x = [u.flatten(2).transpose(1, 2) for u in clean_x] + + seq_lens_clean = torch.tensor([u.size(1) for u in clean_x], dtype=torch.long) + clean_x = torch.cat([ + torch.cat([u, u.new_zeros(1, seq_lens_clean[0] - u.size(1), u.size(2))], dim=1) for u in clean_x + ]) + + x = torch.cat([clean_x, x], dim=1) + if aug_t is None: + aug_t = torch.zeros(t_ori_shape, device=t.device, dtype=t.dtype) + bt_clean = aug_t.size(0) + t_clean_len = aug_t.size(1) + aug_t = aug_t.flatten() + e_clean = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, aug_t).unflatten(0, (bt_clean, t_clean_len)).type_as(x)) + e0_clean = self.time_projection(e_clean).unflatten(2, (6, self.dim)) + e0 = torch.cat([e0_clean, e0], dim=1) + + + # arguments + kwargs = dict( + e=e0, + seq_lens=seq_lens, + grid_sizes=grid_sizes, + freqs=self.freqs, + context=context, + context_lens=context_lens, + block_mask=self.block_mask, + t_scale=self.t_scale, + use_relative_rope=getattr(self, "use_relative_rope", False), + method=getattr(self, "rope_method", "linear"), + original_seq_len=getattr(self, "original_seq_len", None), + temporal_offset=getattr(self, "rope_temporal_offset", 0.0), + ) + + def create_custom_forward(module): + def custom_forward(*inputs, **kwargs): + return module(*inputs, **kwargs) + return custom_forward + + for block in self.blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, **kwargs, + use_reentrant=False, + ) + else: + x = block(x, **kwargs) + + if clean_x is not None: + x = x[:, x.shape[1] // 2:] + + # head + x = self.head(x, e.unsqueeze(2)) + + x = self.unpatchify(x, grid_sizes) + return torch.stack(x) + + +def sp_causal_attn_forward( + self, + x, + seq_lens, + grid_sizes, + freqs, + block_mask, + kv_cache=None, + current_start=0, + cache_start=None, + t_scale=1.0, + use_relative_rope=False, + method="linear", + original_seq_len=None, + temporal_offset=0.0, + **kwargs, +): + r""" + Args: + x(Tensor): Shape [B, L, num_heads, C / num_heads] + seq_lens(Tensor): Shape [B] + grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + block_mask (BlockMask) + t_scale (float): Temporal RoPE interpolation scale. <1.0 compresses positions. + method (str): RoPE method. This release supports "linear". + original_seq_len (int): Unused by the release linear RoPE path. + """ + + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + if cache_start is None: + cache_start = current_start + + # query, key, value function + def qkv_fn(x): + q = self.norm_q(self.q(x)).view(b, s, n, d) + k = self.norm_k(self.k(x)).view(b, s, n, d) + v = self.v(x).view(b, s, n, d) + return q, k, v + + q, k, v = qkv_fn(x) + + if kv_cache is None: + # Teacher-forcing training doubles sequence length with clean/noisy halves. + is_tf = (s == seq_lens[0].item() * 2) + if is_tf: + q_chunk = torch.chunk(q, 2, dim=1) + k_chunk = torch.chunk(k, 2, dim=1) + roped_query = [] + roped_key = [] + # rope should be same for clean and noisy parts + for ii in range(2): + rq = sp_rope_apply(q_chunk[ii], grid_sizes, freqs, t_scale=t_scale, + method=method, original_seq_len=original_seq_len, + temporal_offset=temporal_offset).type_as(v) + rk = sp_rope_apply(k_chunk[ii], grid_sizes, freqs, t_scale=t_scale, + method=method, original_seq_len=original_seq_len, + temporal_offset=temporal_offset).type_as(v) + roped_query.append(rq) + roped_key.append(rk) + + roped_query = torch.cat(roped_query, dim=1) + roped_key = torch.cat(roped_key, dim=1) + + x = distributed_flex_attention( + roped_query, + roped_key, + v, + block_mask, + ) + + else: + roped_query = sp_rope_apply(q, grid_sizes, freqs, t_scale=t_scale, + method=method, original_seq_len=original_seq_len, + temporal_offset=temporal_offset).type_as(v) + roped_key = sp_rope_apply(k, grid_sizes, freqs, t_scale=t_scale, + method=method, original_seq_len=original_seq_len, + temporal_offset=temporal_offset).type_as(v) + + x = distributed_flex_attention( + roped_query, + roped_key, + v, + block_mask, + ) + + else: + raise NotImplementedError() + + # output + x = x.flatten(2) + x = self.o(x) + + # Return both output and cache update info + if kv_cache is not None: + raise NotImplementedError() + else: + return x diff --git a/wan_5b/distributed/sp_training.py b/wan_5b/distributed/sp_training.py new file mode 100644 index 0000000..e65e19f --- /dev/null +++ b/wan_5b/distributed/sp_training.py @@ -0,0 +1,575 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0 +# +# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied. +# +# SPDX-License-Identifier: Apache-2.0 + +"""LongLive sequence-parallel training utilities. + +This module collects the LongLive-specific SP training pieces that are layered +on top of Wan2.2's distributed Ulysses helpers: SP/DP process-group routing, +autograd-safe all-to-all for FlexAttention, distributed teacher-forcing +FlexAttention, and chunk-halo VAE preparation. +""" + +import math + +import torch +import torch.distributed as dist +try: + from torch.nn.attention.flex_attention import flex_attention as _flex_attention +except ModuleNotFoundError: + _flex_attention = None + +from utils.config import wan_default_config + + +DEFAULT_SP_VAE_HALO_LATENTS = 28 + +_sp_group = None +_dp_group = None +_compiled_flex_attention = None + + +def sp_training_sequence_frame_count(config): + """Frames that are sharded by training sequence parallelism.""" + return int(list(config.image_or_video_shape)[1]) + + +def validate_sequence_parallel_training_config(config, sp_size, num_frame_per_block): + total_frames = int(list(config.image_or_video_shape)[1]) + train_frames = sp_training_sequence_frame_count(config) + if train_frames <= 0: + raise ValueError( + f"image_or_video_shape[1] ({total_frames}) must be positive." + ) + if train_frames % int(num_frame_per_block) != 0: + raise ValueError( + f"training latent frames ({train_frames}) must be divisible by " + f"num_frame_per_block ({num_frame_per_block})." + ) + if train_frames % (int(sp_size) * int(num_frame_per_block)) != 0: + raise ValueError( + f"training latent frames ({train_frames}) must be divisible " + f"by sequence_parallel_size ({sp_size}) * num_frame_per_block " + f"({num_frame_per_block})." + ) + + +def _get_compiled_flex_attention(): + global _compiled_flex_attention + if _flex_attention is None: + raise ModuleNotFoundError( + "torch.nn.attention.flex_attention is required for Sequence Parallel " + "training. Install a PyTorch build that provides FlexAttention." + ) + if _compiled_flex_attention is None: + _compiled_flex_attention = torch.compile( + _flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs") + return _compiled_flex_attention + + +def set_sequence_parallel_group(group): + """Set the SP group used by SP rank/world-size and all-to-all helpers.""" + global _sp_group + _sp_group = group + + +def get_sequence_parallel_group(): + return _sp_group + + +def set_data_parallel_group(group): + """Set the DP group whose ranks own the same sequence chunk.""" + global _dp_group + _dp_group = group + + +def get_data_parallel_group(): + return _dp_group + + +def get_sp_rank(): + if _sp_group is not None: + return dist.get_rank(_sp_group) + return dist.get_rank() + + +def get_sp_world_size(): + if _sp_group is not None: + return dist.get_world_size(_sp_group) + return dist.get_world_size() + + +def resolve_sequence_parallel_group(group=None): + if group is not None: + return group + return _sp_group + + +def _world_size(group=None): + return dist.get_world_size(group) if group is not None else dist.get_world_size() + + +def _all_to_all_list_impl(x, scatter_dim, gather_dim, group=None, **kwargs): + world_size = _world_size(group) + if world_size <= 1: + return x + inputs = [u.contiguous() for u in x.chunk(world_size, dim=scatter_dim)] + outputs = [torch.empty_like(u) for u in inputs] + dist.all_to_all(outputs, inputs, group=group, **kwargs) + return torch.cat(outputs, dim=gather_dim).contiguous() + + +def all_to_all(x, scatter_dim, gather_dim, group=None, **kwargs): + group = resolve_sequence_parallel_group(group) + if _world_size(group) <= 1: + return x + return _all_to_all_list_impl(x, scatter_dim, gather_dim, group=group, **kwargs) + + +class AllToAllWithGrad(torch.autograd.Function): + @staticmethod + def forward(ctx, input_tensor, scatter_dim, gather_dim, group): + ctx.scatter_dim = scatter_dim + ctx.gather_dim = gather_dim + ctx.group = group + return _all_to_all_list_impl( + input_tensor, scatter_dim, gather_dim, group=group + ) + + @staticmethod + def backward(ctx, grad_output): + return ( + _all_to_all_list_impl( + grad_output, ctx.gather_dim, ctx.scatter_dim, group=ctx.group + ), + None, + None, + None, + ) + + +def all_to_all_with_grad(x, scatter_dim, gather_dim, group=None): + group = resolve_sequence_parallel_group(group) + if _world_size(group) <= 1: + return x + return AllToAllWithGrad.apply(x, scatter_dim, gather_dim, group) + + +def all_gather(tensor, group=None): + group = resolve_sequence_parallel_group(group) + world_size = _world_size(group) + if world_size == 1: + return [tensor] + tensor_list = [torch.empty_like(tensor) for _ in range(world_size)] + dist.all_gather(tensor_list, tensor, group=group) + return tensor_list + + +def gather_forward(input, dim, group=None): + group = resolve_sequence_parallel_group(group) + if _world_size(group) == 1: + return input + return torch.cat(all_gather(input, group=group), dim=dim).contiguous() + + +def distributed_flex_attention( + roped_q: torch.Tensor, + roped_k: torch.Tensor, + v: torch.Tensor, + block_mask, + pad_multiple: int = 128, +): + """Distributed FlexAttention over SP-sharded heads/sequences. + + Inputs use the local SP layout [B, L_local, N_local, D]. The function + gathers sequence tokens across the SP group, applies block-sparse + FlexAttention, then scatters the output back to the original local layout. + """ + roped_q = all_to_all_with_grad(roped_q, scatter_dim=2, gather_dim=1) + roped_k = all_to_all_with_grad(roped_k, scatter_dim=2, gather_dim=1) + v = all_to_all_with_grad(v, scatter_dim=2, gather_dim=1) + + batch_size, global_len, num_local_heads, head_dim = roped_q.shape + padded_length = math.ceil(global_len / pad_multiple) * pad_multiple - global_len + if padded_length > 0: + zeros_q = torch.zeros( + batch_size, + padded_length, + num_local_heads, + head_dim, + device=roped_q.device, + dtype=v.dtype, + ) + zeros_k = torch.zeros( + batch_size, + padded_length, + num_local_heads, + head_dim, + device=roped_k.device, + dtype=v.dtype, + ) + zeros_v = torch.zeros( + batch_size, + padded_length, + num_local_heads, + head_dim, + device=v.device, + dtype=v.dtype, + ) + roped_q = torch.cat([roped_q, zeros_q], dim=1) + roped_k = torch.cat([roped_k, zeros_k], dim=1) + v = torch.cat([v, zeros_v], dim=1) + + x = _get_compiled_flex_attention()( + query=roped_q.transpose(2, 1), + key=roped_k.transpose(2, 1), + value=v.transpose(2, 1), + block_mask=block_mask, + ) + if padded_length > 0: + x = x[:, :, :-padded_length] + + x = x.transpose(2, 1) + return all_to_all_with_grad(x, scatter_dim=1, gather_dim=2) + + +class SequenceParallelHelper: + def __init__(self, trainer): + self.trainer = trainer + + @property + def config(self): + return self.trainer.config + + @property + def device(self): + return self.trainer.device + + @property + def dtype(self): + return self.trainer.dtype + + @property + def sp_size(self): + return int(getattr(self.trainer, "sequence_parallel_size", 1)) + + @property + def sp_group(self): + return getattr(self.trainer, "sp_group", None) + + @property + def vae_halo_latents(self): + return int(getattr(self.config, "vae_halo_latents", DEFAULT_SP_VAE_HALO_LATENTS)) + + def enabled(self): + return self.sp_size > 1 and self.sp_group is not None + + def sp_root_global_rank(self): + global_rank = dist.get_rank() + return (global_rank // self.sp_size) * self.sp_size + + def local_sp_rank(self): + return get_sp_rank() + + def _chunk_tensor(self, tensor, dim): + return tensor.chunk(self.sp_size, dim=dim)[self.local_sp_rank()].contiguous() + + def local_i2v_initial_latent(self, initial_latent): + if initial_latent is None: + return None + if not getattr(self.config, "i2v", False): + return initial_latent + if not self.enabled(): + return initial_latent + return initial_latent if self.local_sp_rank() == 0 else None + + def partition_training_inputs( + self, + *, + image_or_video_shape, + clean_latent=None, + conditional_dict=None, + clean_latent_is_sharded=False, + ): + if not self.enabled(): + return clean_latent, conditional_dict, image_or_video_shape + + image_or_video_shape = list(image_or_video_shape) + batch_size = ( + clean_latent.shape[0] + if clean_latent is not None + else int(image_or_video_shape[0]) + ) + + def chunk_dict(d): + if d is None: + return None + out = {} + for key, value in d.items(): + if value is None or not torch.is_tensor(value): + out[key] = value + continue + value = value.reshape(batch_size, -1, *value.shape[1:]) + num_segments = value.shape[1] + if num_segments == 1: + out[key] = value + elif num_segments % self.sp_size == 0: + out[key] = self._chunk_tensor(value, dim=1) + else: + raise ValueError( + f"SP chunking failed for key={key}: " + f"num_segments={num_segments} must be 1 or divisible by " + f"sp_size={self.sp_size}" + ) + out[key] = out[key].reshape(-1, *value.shape[2:]).contiguous() + return out + + if clean_latent is not None and not clean_latent_is_sharded: + clean_latent = self._chunk_tensor(clean_latent, dim=1) + + conditional_dict = chunk_dict(conditional_dict) + image_or_video_shape[1] = image_or_video_shape[1] // self.sp_size + return clean_latent, conditional_dict, image_or_video_shape + + def chunk_if_needed(self, tensor, *, dim, already_sharded=False): + if tensor is None or not self.enabled() or already_sharded: + return tensor + return self._chunk_tensor(tensor, dim=dim) + + def build_loss_mask(self, batch, clean_latent, clean_latent_is_sharded): + num_valid_latent_frames = batch.get("num_valid_latent_frames", None) + mask_i2v_first_frame = bool( + getattr(self.config, "i2v", False) + and getattr(self.config, "independent_first_frame", False) + ) + if num_valid_latent_frames is None and not mask_i2v_first_frame: + return None + _, local_or_global_frames = clean_latent.shape[:2] + if clean_latent_is_sharded: + frame_start = self.local_sp_rank() * local_or_global_frames + frame_indices = torch.arange( + frame_start, + frame_start + local_or_global_frames, + device=self.device, + ).unsqueeze(0) + else: + frame_indices = torch.arange( + local_or_global_frames, device=self.device + ).unsqueeze(0) + if num_valid_latent_frames is None: + mask = torch.ones( + (clean_latent.shape[0], local_or_global_frames), + device=self.device, + dtype=torch.float32, + ) + else: + num_valid_latent_frames = num_valid_latent_frames.to(device=self.device) + mask = (frame_indices < num_valid_latent_frames.unsqueeze(1)).float() + if mask_i2v_first_frame: + mask = mask * (frame_indices != 0).float() + return mask + + def partition_loss_mask(self, loss_mask, *, already_sharded=False): + if loss_mask is None: + return None, None + if self.enabled() and not already_sharded: + loss_mask = self._chunk_tensor(loss_mask, dim=1) + if not self.enabled(): + return loss_mask, None + global_valid = loss_mask.sum() + dist.all_reduce(global_valid, op=dist.ReduceOp.SUM, group=self.sp_group) + return loss_mask, global_valid + + def latent_range_to_raw_window(self, latent_start, latent_end): + assert latent_end > latent_start + ratio = int( + wan_default_config[self.config.model_kwargs.model_name][ + "temporal_compression_ratio" + ] + ) + if latent_start == 0: + raw_start = 0 + pseudo_prefix_latents = 0 + else: + raw_start = ratio * (latent_start - 1) + pseudo_prefix_latents = 1 + raw_end = 1 + ratio * (latent_end - 1) + return raw_start, raw_end, pseudo_prefix_latents + + def chunk_halo_meta(self, *, sp_rank, total_latent_frames, total_raw_frames): + if total_latent_frames % self.sp_size != 0: + raise ValueError( + f"total_latent_frames={total_latent_frames} must be divisible by " + f"sequence_parallel_size={self.sp_size}" + ) + local_latent_frames = total_latent_frames // self.sp_size + keep_start = sp_rank * local_latent_frames + keep_end = keep_start + local_latent_frames + halo_start = max(0, keep_start - self.vae_halo_latents) + raw_start, raw_end, pseudo_prefix_latents = self.latent_range_to_raw_window( + halo_start, keep_end + ) + raw_start = max(0, raw_start) + raw_end = min(int(total_raw_frames), raw_end) + drop_latents = pseudo_prefix_latents + (keep_start - halo_start) + return { + "sp_rank": int(sp_rank), + "keep_start": int(keep_start), + "keep_end": int(keep_end), + "halo_start": int(halo_start), + "raw_start": int(raw_start), + "raw_end": int(raw_end), + "raw_frames": int(raw_end - raw_start), + "drop_latents": int(drop_latents), + "local_latent_frames": int(local_latent_frames), + } + + def chunk_halo_metas(self, *, total_latent_frames, total_raw_frames): + return [ + self.chunk_halo_meta( + sp_rank=sp_rank, + total_latent_frames=total_latent_frames, + total_raw_frames=total_raw_frames, + ) + for sp_rank in range(self.sp_size) + ] + + def scatter_frame_windows_for_chunk_halo(self, root_frames): + global_rank = dist.get_rank() + sp_rank = global_rank % self.sp_size + root_global_rank = self.sp_root_global_rank() + + batch_size, channels, total_raw_frames, height, width = tuple(root_frames.shape) + total_latent_frames = int(list(self.config.image_or_video_shape)[1]) + metas = self.chunk_halo_metas( + total_latent_frames=total_latent_frames, + total_raw_frames=total_raw_frames, + ) + rank_meta = metas[sp_rank] + + if global_rank == root_global_rank: + local_frames = None + for send_meta in metas: + raw_window = root_frames[ + :, + :, + send_meta["raw_start"]:send_meta["raw_end"], + :, + :, + ] + dst_global_rank = root_global_rank + int(send_meta["sp_rank"]) + if dst_global_rank == global_rank: + local_frames = raw_window.contiguous() + continue + send_buffer = raw_window.contiguous() + dist.send(send_buffer, dst=dst_global_rank) + del send_buffer + if local_frames is None: + raise RuntimeError("SP-VAE chunk_halo root did not build local frames") + else: + local_frames = torch.empty( + (batch_size, channels, rank_meta["raw_frames"], height, width), + device=self.device, + dtype=root_frames.dtype, + ) + dist.recv(local_frames, src=root_global_rank) + local_frames = local_frames.contiguous() + + return local_frames, rank_meta + + def sync_batch(self, batch, step): + if not self.enabled(): + return batch + + global_rank = dist.get_rank() + root_global_rank = self.sp_root_global_rank() + device = torch.device(f"cuda:{torch.cuda.current_device()}") + + if "frames" in batch: + frames_for_scatter = batch["frames"] + if global_rank == root_global_rank: + frames_for_scatter = frames_for_scatter.to(device, non_blocking=True) + local_frames, meta = self.scatter_frame_windows_for_chunk_halo( + frames_for_scatter + ) + batch["frames"] = local_frames + batch["sp_vae_chunk_meta"] = meta + + if global_rank == root_global_rank: + info_obj = [{ + "prompts": batch["prompts"], + "idx": batch["idx"], + "num_valid_latent_frames": batch.get("num_valid_latent_frames", None), + }] + else: + info_obj = [None] + dist.broadcast_object_list(info_obj, src=root_global_rank, group=self.sp_group) + + batch["prompts"] = info_obj[0]["prompts"] + batch["idx"] = info_obj[0]["idx"] + if info_obj[0]["num_valid_latent_frames"] is not None: + batch["num_valid_latent_frames"] = info_obj[0]["num_valid_latent_frames"] + return batch + + def broadcast_tensor_from_root(self, root_tensor, *, shape): + root_global_rank = self.sp_root_global_rank() + if dist.get_rank() == root_global_rank: + tensor = root_tensor.contiguous() + else: + tensor = torch.empty(shape, device=self.device, dtype=self.dtype) + dist.broadcast(tensor, src=root_global_rank, group=self.sp_group) + return tensor + + def encode_raw_video_latents(self, batch, *, batch_size): + if self.enabled(): + meta = batch.get("sp_vae_chunk_meta", None) + if meta is None: + raise RuntimeError( + "SP chunk-halo VAE requires sync_batch to attach sp_vae_chunk_meta." + ) + total_latent_frames = int(list(self.config.image_or_video_shape)[1]) + + with torch.no_grad(): + frames = batch["frames"].to( + device=self.device, dtype=self.dtype, non_blocking=True + ) + chunk_latent = self.trainer.model.vae.encode_to_latent(frames).to( + device=self.device, dtype=self.dtype + ) + + drop = int(meta["drop_latents"]) + local_latent_frames = int(meta["local_latent_frames"]) + keep_end = drop + local_latent_frames + if chunk_latent.shape[1] < keep_end: + raise RuntimeError( + "SP-VAE chunk_halo encode produced too few latent frames: " + f"chunk_latent={tuple(chunk_latent.shape)} drop={drop} " + f"local_latent_frames={local_latent_frames} meta={meta}" + ) + clean_latent = chunk_latent[:, drop:keep_end].contiguous() + if clean_latent.shape[1] != total_latent_frames // self.sp_size: + raise RuntimeError( + "SP-VAE chunk_halo local latent shape mismatch: " + f"clean_latent={tuple(clean_latent.shape)} meta={meta}" + ) + + image_latent = ( + clean_latent[:, 0:1] if int(meta["keep_start"]) == 0 else None + ) + + del chunk_latent + return clean_latent, image_latent, True + + frames = batch["frames"].to(device=self.device, dtype=self.dtype) + with torch.no_grad(): + clean_latent = self.trainer.model.vae.encode_to_latent(frames).to( + device=self.device, dtype=self.dtype + ) + image_latent = clean_latent[:, 0:1] + return clean_latent, image_latent, False diff --git a/wan_5b/distributed/sp_ulysses_inference.py b/wan_5b/distributed/sp_ulysses_inference.py new file mode 100644 index 0000000..8349d8c --- /dev/null +++ b/wan_5b/distributed/sp_ulysses_inference.py @@ -0,0 +1,209 @@ +# Copyright 2024-2025 LongLive Authors. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +""" +Ulysses-style sequence-parallel primitives for Wan2.2-TI2V-5B inference. +""" + +from typing import Optional +import time + +import torch +import torch.distributed as dist + + +_SP_GROUP: Optional[dist.ProcessGroup] = None +_SP_WORLD_SIZE: int = 1 +_SP_RANK: int = 0 + +_SP_COMM_STATS = { + "all_gather_time": 0.0, + "all_gather_count": 0, + "all_gather_bytes": 0, + "all_to_all_time": 0.0, + "all_to_all_count": 0, + "all_to_all_bytes": 0, + "barrier_time": 0.0, + "barrier_count": 0, +} +_SP_PROFILING_ENABLED: bool = False + + +def init_sequence_parallel( + world_size: int | None = None, + rank: int | None = None, + group: Optional[dist.ProcessGroup] = None, +): + """Initialize sequence-parallel state for the current process.""" + global _SP_GROUP, _SP_WORLD_SIZE, _SP_RANK + if not dist.is_initialized(): + raise RuntimeError("torch.distributed must be initialized before init_sequence_parallel") + if group is not None: + _SP_GROUP = group + _SP_WORLD_SIZE = dist.get_world_size(group) if world_size is None else world_size + _SP_RANK = dist.get_rank(group) if rank is None else rank + else: + _SP_WORLD_SIZE = world_size if world_size is not None else dist.get_world_size() + _SP_RANK = rank if rank is not None else dist.get_rank() + _SP_GROUP = dist.group.WORLD + if _SP_RANK == 0: + print(f"[SP-Ulysses-5B] Initialized: world_size={_SP_WORLD_SIZE}, rank={_SP_RANK}") + return _SP_GROUP + + +def get_sp_group() -> dist.ProcessGroup: + if _SP_GROUP is None: + raise RuntimeError("SP group not initialized. Call init_sequence_parallel first.") + return _SP_GROUP + + +def get_sp_world_size() -> int: + return _SP_WORLD_SIZE + + +def get_sp_rank() -> int: + return _SP_RANK + + +def is_sp_enabled() -> bool: + return _SP_WORLD_SIZE > 1 + + +def enable_sp_profiling(): + global _SP_PROFILING_ENABLED + _SP_PROFILING_ENABLED = True + reset_sp_comm_stats() + + +def disable_sp_profiling(): + global _SP_PROFILING_ENABLED + _SP_PROFILING_ENABLED = False + + +def reset_sp_comm_stats(): + global _SP_COMM_STATS + _SP_COMM_STATS = { + key: 0.0 if "time" in key or "bytes" in key else 0 + for key in _SP_COMM_STATS + } + + +def get_sp_comm_stats(): + return _SP_COMM_STATS.copy() + + +def sp_all_gather(tensor: torch.Tensor, dim: int = 1) -> torch.Tensor: + if not is_sp_enabled(): + return tensor + global _SP_COMM_STATS, _SP_PROFILING_ENABLED + world_size = get_sp_world_size() + tensor_list = [torch.empty_like(tensor) for _ in range(world_size)] + if _SP_PROFILING_ENABLED: + torch.cuda.synchronize() + start_time = time.perf_counter() + dist.all_gather(tensor_list, tensor, group=get_sp_group()) + if _SP_PROFILING_ENABLED: + torch.cuda.synchronize() + elapsed = time.perf_counter() - start_time + _SP_COMM_STATS["all_gather_time"] += elapsed + _SP_COMM_STATS["all_gather_count"] += 1 + _SP_COMM_STATS["all_gather_bytes"] += ( + tensor.numel() * tensor.element_size() * (world_size - 1) + ) + return torch.cat(tensor_list, dim=dim) + + +def sp_scatter(tensor: torch.Tensor, dim: int = 1) -> torch.Tensor: + if not is_sp_enabled(): + return tensor + chunks = torch.chunk(tensor, get_sp_world_size(), dim=dim) + return chunks[get_sp_rank()].contiguous() + + +def sp_all_to_all(tensor: torch.Tensor, scatter_dim: int, gather_dim: int) -> torch.Tensor: + if not is_sp_enabled(): + return tensor + global _SP_COMM_STATS, _SP_PROFILING_ENABLED + world_size = get_sp_world_size() + if _SP_PROFILING_ENABLED: + torch.cuda.synchronize() + start_time = time.perf_counter() + scatter_chunks = [ + chunk.contiguous() for chunk in torch.chunk(tensor, world_size, dim=scatter_dim) + ] + recv_chunks = [torch.empty_like(scatter_chunks[0]) for _ in range(world_size)] + dist.all_to_all(recv_chunks, scatter_chunks, group=get_sp_group()) + output = torch.cat(recv_chunks, dim=gather_dim) + if _SP_PROFILING_ENABLED: + torch.cuda.synchronize() + elapsed = time.perf_counter() - start_time + _SP_COMM_STATS["all_to_all_time"] += elapsed + _SP_COMM_STATS["all_to_all_count"] += 1 + _SP_COMM_STATS["all_to_all_bytes"] += ( + scatter_chunks[0].numel() * tensor.element_size() * (world_size - 1) * 2 + ) + return output + + +def ulysses_seq_to_head(tensor: torch.Tensor) -> torch.Tensor: + """Convert [B, S_local, N, D] to [B, S_total, N_local, D].""" + return sp_all_to_all(tensor, scatter_dim=2, gather_dim=1) if is_sp_enabled() else tensor + + +def ulysses_head_to_seq(tensor: torch.Tensor) -> torch.Tensor: + """Convert [B, S_total, N_local, D] to [B, S_local, N, D].""" + return sp_all_to_all(tensor, scatter_dim=1, gather_dim=2) if is_sp_enabled() else tensor + + +def sp_barrier(): + if not is_sp_enabled(): + return + global _SP_COMM_STATS, _SP_PROFILING_ENABLED + if _SP_PROFILING_ENABLED: + torch.cuda.synchronize() + start_time = time.perf_counter() + dist.barrier(group=get_sp_group()) + if _SP_PROFILING_ENABLED: + torch.cuda.synchronize() + elapsed = time.perf_counter() - start_time + _SP_COMM_STATS["barrier_time"] += elapsed + _SP_COMM_STATS["barrier_count"] += 1 + + +def sp_print(msg: str, rank_only: int = 0): + if get_sp_rank() == rank_only: + print(f"[SP-5B Rank {get_sp_rank()}] {msg}") + + +def profile_sp_communication(): + if not is_sp_enabled(): + return + rank = get_sp_rank() + world_size = get_sp_world_size() + test_size = (1, 880, 24, 128) + test_tensor = torch.randn(test_size, device="cuda", dtype=torch.bfloat16) + for _ in range(3): + _ = sp_all_gather(test_tensor, dim=1) + _ = ulysses_seq_to_head(test_tensor) + torch.cuda.synchronize() + start = time.perf_counter() + for _ in range(10): + _ = sp_all_gather(test_tensor, dim=1) + torch.cuda.synchronize() + all_gather_time = (time.perf_counter() - start) / 10 * 1000 + start = time.perf_counter() + for _ in range(10): + _ = ulysses_seq_to_head(test_tensor) + torch.cuda.synchronize() + all_to_all_time = (time.perf_counter() - start) / 10 * 1000 + if rank == 0: + all_gather_bw = ( + test_tensor.numel() * test_tensor.element_size() * (world_size - 1) / 1e9 + ) / (all_gather_time / 1000) + all_to_all_bw = ( + test_tensor.numel() * test_tensor.element_size() * (world_size - 1) / 1e9 + ) / (all_to_all_time / 1000) + print("\n[SP-Ulysses-5B Profile]") + print(f" World Size: {world_size}") + print(f" Test Shape: {test_size}") + print(f" All-Gather: {all_gather_time:.2f} ms, Bandwidth: {all_gather_bw:.2f} GB/s") + print(f" All-to-All: {all_to_all_time:.2f} ms, Bandwidth: {all_to_all_bw:.2f} GB/s") diff --git a/wan_5b/distributed/ulysses.py b/wan_5b/distributed/ulysses.py new file mode 100644 index 0000000..a6c229c --- /dev/null +++ b/wan_5b/distributed/ulysses.py @@ -0,0 +1,47 @@ +# Adopted from https://github.com/Wan-Video/Wan2.2 +# SPDX-License-Identifier: Apache-2.0 + +import torch.distributed as dist + +from ..modules.attention import flash_attention +from .sp_training import distributed_flex_attention +from .util import all_to_all + + +def distributed_attention( + q, + k, + v, + seq_lens, + window_size=(-1, -1), +): + """ + Performs distributed attention based on DeepSpeed Ulysses attention mechanism. + please refer to https://arxiv.org/pdf/2309.14509 + + Args: + q: [B, Lq // p, Nq, C1]. + k: [B, Lk // p, Nk, C1]. + v: [B, Lk // p, Nk, C2]. Nq must be divisible by Nk. + seq_lens: [B], length of each sequence in batch + window_size: (left right). If not (-1, -1), apply sliding window local attention. + """ + if not dist.is_initialized(): + raise ValueError("distributed group should be initialized.") + + q = all_to_all(q, scatter_dim=2, gather_dim=1) + k = all_to_all(k, scatter_dim=2, gather_dim=1) + v = all_to_all(v, scatter_dim=2, gather_dim=1) + + x = flash_attention( + q, + k, + v, + k_lens=seq_lens, + window_size=window_size, + ) + + return all_to_all(x, scatter_dim=1, gather_dim=2) + + +__all__ = ["distributed_attention", "distributed_flex_attention"] diff --git a/wan_5b/distributed/util.py b/wan_5b/distributed/util.py new file mode 100644 index 0000000..e33f497 --- /dev/null +++ b/wan_5b/distributed/util.py @@ -0,0 +1,51 @@ +# Adopted from https://github.com/Wan-Video/Wan2.2 +# SPDX-License-Identifier: Apache-2.0 + +"""Wan distributed utility compatibility layer. + +LongLive-specific SP/DP group routing and autograd all-to-all live in +``wan_5b.distributed.sp_training``. This module keeps Wan2.2's public import +paths intact for the rest of the codebase. +""" + +import torch.distributed as dist + +from .sp_training import ( + all_gather, + all_to_all, + all_to_all_with_grad, + gather_forward, + get_data_parallel_group, + get_sp_rank, + get_sp_world_size, + set_data_parallel_group, + set_sequence_parallel_group, +) + + +def init_distributed_group(): + """Initialize the default distributed group when it is not yet ready.""" + if not dist.is_initialized(): + dist.init_process_group(backend="nccl") + + +def get_rank(): + return get_sp_rank() + + +def get_world_size(): + return get_sp_world_size() + + +__all__ = [ + "all_gather", + "all_to_all", + "all_to_all_with_grad", + "gather_forward", + "get_data_parallel_group", + "get_rank", + "get_world_size", + "init_distributed_group", + "set_data_parallel_group", + "set_sequence_parallel_group", +] diff --git a/wan_5b/image2video.py b/wan_5b/image2video.py new file mode 100644 index 0000000..659564c --- /dev/null +++ b/wan_5b/image2video.py @@ -0,0 +1,431 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import gc +import logging +import math +import os +import random +import sys +import types +from contextlib import contextmanager +from functools import partial + +import numpy as np +import torch +import torch.cuda.amp as amp +import torch.distributed as dist +import torchvision.transforms.functional as TF +from tqdm import tqdm + +from .distributed.fsdp import shard_model +from .distributed.sequence_parallel import sp_attn_forward, sp_dit_forward +from .distributed.util import get_world_size +from .modules.model import WanModel +from .modules.t5 import T5EncoderModel +from .modules.vae2_1 import Wan2_1_VAE +from .utils.fm_solvers import ( + FlowDPMSolverMultistepScheduler, + get_sampling_sigmas, + retrieve_timesteps, +) +from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler + + +class WanI2V: + + def __init__( + self, + config, + checkpoint_dir, + device_id=0, + rank=0, + t5_fsdp=False, + dit_fsdp=False, + use_sp=False, + t5_cpu=False, + init_on_cpu=True, + convert_model_dtype=False, + ): + r""" + Initializes the image-to-video generation model components. + + Args: + config (EasyDict): + Object containing model parameters initialized from config.py + checkpoint_dir (`str`): + Path to directory containing model checkpoints + device_id (`int`, *optional*, defaults to 0): + Id of target GPU device + rank (`int`, *optional*, defaults to 0): + Process rank for distributed training + t5_fsdp (`bool`, *optional*, defaults to False): + Enable FSDP sharding for T5 model + dit_fsdp (`bool`, *optional*, defaults to False): + Enable FSDP sharding for DiT model + use_sp (`bool`, *optional*, defaults to False): + Enable distribution strategy of sequence parallel. + t5_cpu (`bool`, *optional*, defaults to False): + Whether to place T5 model on CPU. Only works without t5_fsdp. + init_on_cpu (`bool`, *optional*, defaults to True): + Enable initializing Transformer Model on CPU. Only works without FSDP or USP. + convert_model_dtype (`bool`, *optional*, defaults to False): + Convert DiT model parameters dtype to 'config.param_dtype'. + Only works without FSDP. + """ + self.device = torch.device(f"cuda:{device_id}") + self.config = config + self.rank = rank + self.t5_cpu = t5_cpu + self.init_on_cpu = init_on_cpu + + self.num_train_timesteps = config.num_train_timesteps + self.boundary = config.boundary + self.param_dtype = config.param_dtype + + if t5_fsdp or dit_fsdp or use_sp: + self.init_on_cpu = False + + shard_fn = partial(shard_model, device_id=device_id) + self.text_encoder = T5EncoderModel( + text_len=config.text_len, + dtype=config.t5_dtype, + device=torch.device('cpu'), + checkpoint_path=os.path.join(checkpoint_dir, config.t5_checkpoint), + tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer), + shard_fn=shard_fn if t5_fsdp else None, + ) + + self.vae_stride = config.vae_stride + self.patch_size = config.patch_size + self.vae = Wan2_1_VAE( + vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint), + device=self.device) + + logging.info(f"Creating WanModel from {checkpoint_dir}") + self.low_noise_model = WanModel.from_pretrained( + checkpoint_dir, subfolder=config.low_noise_checkpoint) + self.low_noise_model = self._configure_model( + model=self.low_noise_model, + use_sp=use_sp, + dit_fsdp=dit_fsdp, + shard_fn=shard_fn, + convert_model_dtype=convert_model_dtype) + + self.high_noise_model = WanModel.from_pretrained( + checkpoint_dir, subfolder=config.high_noise_checkpoint) + self.high_noise_model = self._configure_model( + model=self.high_noise_model, + use_sp=use_sp, + dit_fsdp=dit_fsdp, + shard_fn=shard_fn, + convert_model_dtype=convert_model_dtype) + if use_sp: + self.sp_size = get_world_size() + else: + self.sp_size = 1 + + self.sample_neg_prompt = config.sample_neg_prompt + + def _configure_model(self, model, use_sp, dit_fsdp, shard_fn, + convert_model_dtype): + """ + Configures a model object. This includes setting evaluation modes, + applying distributed parallel strategy, and handling device placement. + + Args: + model (torch.nn.Module): + The model instance to configure. + use_sp (`bool`): + Enable distribution strategy of sequence parallel. + dit_fsdp (`bool`): + Enable FSDP sharding for DiT model. + shard_fn (callable): + The function to apply FSDP sharding. + convert_model_dtype (`bool`): + Convert DiT model parameters dtype to 'config.param_dtype'. + Only works without FSDP. + + Returns: + torch.nn.Module: + The configured model. + """ + model.eval().requires_grad_(False) + + if use_sp: + for block in model.blocks: + block.self_attn.forward = types.MethodType( + sp_attn_forward, block.self_attn) + model.forward = types.MethodType(sp_dit_forward, model) + + if dist.is_initialized(): + dist.barrier() + + if dit_fsdp: + model = shard_fn(model) + else: + if convert_model_dtype: + model.to(self.param_dtype) + if not self.init_on_cpu: + model.to(self.device) + + return model + + def _prepare_model_for_timestep(self, t, boundary, offload_model): + r""" + Prepares and returns the required model for the current timestep. + + Args: + t (torch.Tensor): + current timestep. + boundary (`int`): + The timestep threshold. If `t` is at or above this value, + the `high_noise_model` is considered as the required model. + offload_model (`bool`): + A flag intended to control the offloading behavior. + + Returns: + torch.nn.Module: + The active model on the target device for the current timestep. + """ + if t.item() >= boundary: + required_model_name = 'high_noise_model' + offload_model_name = 'low_noise_model' + else: + required_model_name = 'low_noise_model' + offload_model_name = 'high_noise_model' + if offload_model or self.init_on_cpu: + if next(getattr( + self, + offload_model_name).parameters()).device.type == 'cuda': + getattr(self, offload_model_name).to('cpu') + if next(getattr( + self, + required_model_name).parameters()).device.type == 'cpu': + getattr(self, required_model_name).to(self.device) + return getattr(self, required_model_name) + + def generate(self, + input_prompt, + img, + max_area=720 * 1280, + frame_num=81, + shift=5.0, + sample_solver='unipc', + sampling_steps=40, + guide_scale=5.0, + n_prompt="", + seed=-1, + offload_model=True): + r""" + Generates video frames from input image and text prompt using diffusion process. + + Args: + input_prompt (`str`): + Text prompt for content generation. + img (PIL.Image.Image): + Input image tensor. Shape: [3, H, W] + max_area (`int`, *optional*, defaults to 720*1280): + Maximum pixel area for latent space calculation. Controls video resolution scaling + frame_num (`int`, *optional*, defaults to 81): + How many frames to sample from a video. The number should be 4n+1 + shift (`float`, *optional*, defaults to 5.0): + Noise schedule shift parameter. Affects temporal dynamics + [NOTE]: If you want to generate a 480p video, it is recommended to set the shift value to 3.0. + sample_solver (`str`, *optional*, defaults to 'unipc'): + Solver used to sample the video. + sampling_steps (`int`, *optional*, defaults to 40): + Number of diffusion sampling steps. Higher values improve quality but slow generation + guide_scale (`float` or tuple[`float`], *optional*, defaults 5.0): + Classifier-free guidance scale. Controls prompt adherence vs. creativity. + If tuple, the first guide_scale will be used for low noise model and + the second guide_scale will be used for high noise model. + n_prompt (`str`, *optional*, defaults to ""): + Negative prompt for content exclusion. If not given, use `config.sample_neg_prompt` + seed (`int`, *optional*, defaults to -1): + Random seed for noise generation. If -1, use random seed + offload_model (`bool`, *optional*, defaults to True): + If True, offloads models to CPU during generation to save VRAM + + Returns: + torch.Tensor: + Generated video frames tensor. Dimensions: (C, N H, W) where: + - C: Color channels (3 for RGB) + - N: Number of frames (81) + - H: Frame height (from max_area) + - W: Frame width from max_area) + """ + # preprocess + guide_scale = (guide_scale, guide_scale) if isinstance( + guide_scale, float) else guide_scale + img = TF.to_tensor(img).sub_(0.5).div_(0.5).to(self.device) + + F = frame_num + h, w = img.shape[1:] + aspect_ratio = h / w + lat_h = round( + np.sqrt(max_area * aspect_ratio) // self.vae_stride[1] // + self.patch_size[1] * self.patch_size[1]) + lat_w = round( + np.sqrt(max_area / aspect_ratio) // self.vae_stride[2] // + self.patch_size[2] * self.patch_size[2]) + h = lat_h * self.vae_stride[1] + w = lat_w * self.vae_stride[2] + + max_seq_len = ((F - 1) // self.vae_stride[0] + 1) * lat_h * lat_w // ( + self.patch_size[1] * self.patch_size[2]) + max_seq_len = int(math.ceil(max_seq_len / self.sp_size)) * self.sp_size + + seed = seed if seed >= 0 else random.randint(0, sys.maxsize) + seed_g = torch.Generator(device=self.device) + seed_g.manual_seed(seed) + noise = torch.randn( + 16, + (F - 1) // self.vae_stride[0] + 1, + lat_h, + lat_w, + dtype=torch.float32, + generator=seed_g, + device=self.device) + + msk = torch.ones(1, F, lat_h, lat_w, device=self.device) + msk[:, 1:] = 0 + msk = torch.concat([ + torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:] + ], + dim=1) + msk = msk.view(1, msk.shape[1] // 4, 4, lat_h, lat_w) + msk = msk.transpose(1, 2)[0] + + if n_prompt == "": + n_prompt = self.sample_neg_prompt + + # preprocess + if not self.t5_cpu: + self.text_encoder.model.to(self.device) + context = self.text_encoder([input_prompt], self.device) + context_null = self.text_encoder([n_prompt], self.device) + if offload_model: + self.text_encoder.model.cpu() + else: + context = self.text_encoder([input_prompt], torch.device('cpu')) + context_null = self.text_encoder([n_prompt], torch.device('cpu')) + context = [t.to(self.device) for t in context] + context_null = [t.to(self.device) for t in context_null] + + y = self.vae.encode([ + torch.concat([ + torch.nn.functional.interpolate( + img[None].cpu(), size=(h, w), mode='bicubic').transpose( + 0, 1), + torch.zeros(3, F - 1, h, w) + ], + dim=1).to(self.device) + ])[0] + y = torch.concat([msk, y]) + + @contextmanager + def noop_no_sync(): + yield + + no_sync_low_noise = getattr(self.low_noise_model, 'no_sync', + noop_no_sync) + no_sync_high_noise = getattr(self.high_noise_model, 'no_sync', + noop_no_sync) + + # evaluation mode + with ( + torch.amp.autocast('cuda', dtype=self.param_dtype), + torch.no_grad(), + no_sync_low_noise(), + no_sync_high_noise(), + ): + boundary = self.boundary * self.num_train_timesteps + + if sample_solver == 'unipc': + sample_scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=self.num_train_timesteps, + shift=1, + use_dynamic_shifting=False) + sample_scheduler.set_timesteps( + sampling_steps, device=self.device, shift=shift) + timesteps = sample_scheduler.timesteps + elif sample_solver == 'dpm++': + sample_scheduler = FlowDPMSolverMultistepScheduler( + num_train_timesteps=self.num_train_timesteps, + shift=1, + use_dynamic_shifting=False) + sampling_sigmas = get_sampling_sigmas(sampling_steps, shift) + timesteps, _ = retrieve_timesteps( + sample_scheduler, + device=self.device, + sigmas=sampling_sigmas) + else: + raise NotImplementedError("Unsupported solver.") + + # sample videos + latent = noise + + arg_c = { + 'context': [context[0]], + 'seq_len': max_seq_len, + 'y': [y], + } + + arg_null = { + 'context': context_null, + 'seq_len': max_seq_len, + 'y': [y], + } + + if offload_model: + torch.cuda.empty_cache() + + for _, t in enumerate(tqdm(timesteps)): + latent_model_input = [latent.to(self.device)] + timestep = [t] + + timestep = torch.stack(timestep).to(self.device) + + model = self._prepare_model_for_timestep( + t, boundary, offload_model) + sample_guide_scale = guide_scale[1] if t.item( + ) >= boundary else guide_scale[0] + + noise_pred_cond = model( + latent_model_input, t=timestep, **arg_c)[0] + if offload_model: + torch.cuda.empty_cache() + noise_pred_uncond = model( + latent_model_input, t=timestep, **arg_null)[0] + if offload_model: + torch.cuda.empty_cache() + noise_pred = noise_pred_uncond + sample_guide_scale * ( + noise_pred_cond - noise_pred_uncond) + + temp_x0 = sample_scheduler.step( + noise_pred.unsqueeze(0), + t, + latent.unsqueeze(0), + return_dict=False, + generator=seed_g)[0] + latent = temp_x0.squeeze(0) + + x0 = [latent] + del latent_model_input, timestep + + if offload_model: + self.low_noise_model.cpu() + self.high_noise_model.cpu() + torch.cuda.empty_cache() + + if self.rank == 0: + videos = self.vae.decode(x0) + + del noise, latent, x0 + del sample_scheduler + if offload_model: + gc.collect() + torch.cuda.synchronize() + if dist.is_initialized(): + dist.barrier() + + return videos[0] if self.rank == 0 else None diff --git a/wan_5b/modules/__init__.py b/wan_5b/modules/__init__.py new file mode 100644 index 0000000..9d9eeb8 --- /dev/null +++ b/wan_5b/modules/__init__.py @@ -0,0 +1,19 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +from .attention import flash_attention +from .model import WanModel +from .t5 import T5Decoder, T5Encoder, T5EncoderModel, T5Model +from .tokenizers import HuggingfaceTokenizer +from .vae2_1 import Wan2_1_VAE +from .vae2_2 import Wan2_2_VAE + +__all__ = [ + 'Wan2_1_VAE', + 'Wan2_2_VAE', + 'WanModel', + 'T5Model', + 'T5Encoder', + 'T5Decoder', + 'T5EncoderModel', + 'HuggingfaceTokenizer', + 'flash_attention', +] diff --git a/wan_5b/modules/attention.py b/wan_5b/modules/attention.py new file mode 100644 index 0000000..b6b8653 --- /dev/null +++ b/wan_5b/modules/attention.py @@ -0,0 +1,306 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import functools +import os +import torch + +# FA4: preferred on Blackwell (sm_100). Imported via `flash_attn.cute`. The +# inference loop on GB200 runs this path when LLV2_USE_FA4 is unset or "1". +try: + from flash_attn.cute import flash_attn_varlen_func as _fa4_varlen_func + FLASH_ATTN_4_AVAILABLE = True +except Exception: + FLASH_ATTN_4_AVAILABLE = False + +try: + import flash_attn_interface + FLASH_ATTN_3_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_3_AVAILABLE = False + +try: + import flash_attn + FLASH_ATTN_2_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_2_AVAILABLE = False + +# TE 2.13 ships a `DotProductAttention` Module whose cuDNN-backed +# FusedAttention path is Blackwell sm_100-tuned. Default it ON when TE is +# importable; couple it with `NVTE_FLASH_ATTN=0` in the launch env so TE +# picks cuDNN FusedAttention instead of falling back to flash-attn 2 (which +# is what the original baseline used and is the kernel we want to replace). +try: + from transformer_engine.pytorch.attention import ( + DotProductAttention as _TE_DPA, + ) + TE_DPA_AVAILABLE = True +except Exception: + TE_DPA_AVAILABLE = False + +# Default off — iter-5 (run-20260520-022256-fa4) showed FA4 4.0.0b4 + torch +# 2.12 + cute-DSL is currently quality-FAIL on this NVFP4 pipeline (max|Δ|≈9 +# vs threshold 5e-3) and ~6% slower in steady-state. Re-enable with +# LLV2_USE_FA4=1 once we have a clean qlive_fa4 baseline + working +# torch.compile interop. +_USE_FA4 = os.environ.get("LLV2_USE_FA4", "0") == "1" +_USE_TE_ATTN = os.environ.get("LLV2_USE_TE_ATTN", "0") == "1" +# iter-32: FA3 default-off. Initial sm_100 build only JIT'd common head_dim +# templates; less-common shapes throw "no kernel image is available". Rebuild +# FA3 with TORCH_CUDA_ARCH_LIST=10.0+PTX then flip this to 1. +_USE_FA3 = os.environ.get("LLV2_USE_FA3", "0") == "1" + + +@functools.lru_cache(maxsize=16) +def _get_te_dpa( + num_heads: int, + head_dim: int, + attn_mask_type: str, + window_left: int, + window_right: int, +) -> "torch.nn.Module": + """Cached TE DotProductAttention instance keyed by attention shape + + masking. Constructed lazily and reused across forward calls. TE's DPA + object is light at __init__ (no params); the cuDNN dispatch happens in + forward. + """ + ws = (window_left, window_right) + return _TE_DPA( + num_attention_heads=num_heads, + kv_channels=head_dim, + attention_dropout=0.0, + attn_mask_type=attn_mask_type, + window_size=ws, + qkv_format="thd", # varlen — flat tokens + cu_seqlens + ).cuda() + + +import warnings + +__all__ = [ + 'flash_attention', + 'attention', +] + + +def flash_attention( + q, + k, + v, + q_lens=None, + k_lens=None, + dropout_p=0., + softmax_scale=None, + q_scale=None, + causal=False, + window_size=(-1, -1), + deterministic=False, + dtype=torch.bfloat16, + version=None, +): + """ + q: [B, Lq, Nq, C1]. + k: [B, Lk, Nk, C1]. + v: [B, Lk, Nk, C2]. Nq must be divisible by Nk. + q_lens: [B]. + k_lens: [B]. + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + causal: bool. Whether to apply causal attention mask. + window_size: (left right). If not (-1, -1), apply sliding window local attention. + deterministic: bool. If True, slightly slower and uses more memory. + dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16. + """ + half_dtypes = (torch.float16, torch.bfloat16) + assert dtype in half_dtypes + assert q.device.type == 'cuda' and q.size(-1) <= 256 + + # params + b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype + + def half(x): + return x if x.dtype in half_dtypes else x.to(dtype) + + # preprocess query + if q_lens is None: + q = half(q.flatten(0, 1)) + q_lens = torch.tensor( + [lq] * b, dtype=torch.int32).to( + device=q.device, non_blocking=True) + else: + q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)])) + + # preprocess key, value + if k_lens is None: + k = half(k.flatten(0, 1)) + v = half(v.flatten(0, 1)) + k_lens = torch.tensor( + [lk] * b, dtype=torch.int32).to( + device=k.device, non_blocking=True) + else: + k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)])) + v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)])) + + q = q.to(v.dtype) + k = k.to(v.dtype) + + if q_scale is not None: + q = q * q_scale + + if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE: + warnings.warn( + 'Flash attention 3 is not available, use flash attention 2 instead.' + ) + + cu_seqlens_q = torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum( + 0, dtype=torch.int32).to(q.device, non_blocking=True) + cu_seqlens_k = torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum( + 0, dtype=torch.int32).to(q.device, non_blocking=True) + + # TE DotProductAttention (cuDNN FusedAttention, sm_100-tuned). Opt-in via + # LLV2_USE_TE_ATTN=1 + NVTE_FLASH_ATTN=0 in the launch env (the latter + # stops TE from dispatching internally to flash-attn 2 — which is what + # we're trying to replace). + # + # iter-6 unit test (agent/te_dpa_unit_test.py, ran in qlive env outside + # any TE FP8 autocast scope) showed `padding`+`window=(-1,-1)` matches + # FA2 (causal=False) at max|Δ|=3e-5 (bf16 rounding). But iter-6 in the + # full pipeline gave video PSNR = 10.4 dB — the math goes wrong because + # the model's TE-wrapped Linear forwards open a `te.fp8_autocast(...)` + # scope, and the DPA inside that scope tries to run FP8 attention without + # calibrated scales. Wrapping the DPA call in `fp8_autocast(enabled=False)` + # forces it to bf16 cuDNN attention, which is what the unit test verified. + if _USE_TE_ATTN and TE_DPA_AVAILABLE: + n_q = q.size(1) # after flatten(0,1), q is [Lq_total, n, d]; size(1)=n + d = q.size(2) + ws_left = -1 if window_size[0] is None or window_size[0] < 0 else int(window_size[0]) + ws_right = -1 if window_size[1] is None or window_size[1] < 0 else int(window_size[1]) + mask_type = "padding_causal" if causal else "padding" + if q_scale is not None and softmax_scale is None: + softmax_scale = float(q_scale) / (d ** 0.5) + dpa = _get_te_dpa(n_q, d, mask_type, ws_left, ws_right) + # iter-6b confirmed wrapping each DPA call in a fp8_autocast(enabled=False) + # context is (a) a no-op for correctness (latent drift unchanged) and + # (b) a recompile trap for dynamo (medians thrash between 1272 and 1859 + # across prompts). Just call DPA directly. + out = dpa( + q, k, v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_k, + max_seqlen_q=lq, + max_seqlen_kv=lk, + ) + if out.dim() == 2: + out = out.view(-1, n_q, d) + x = out.unflatten(0, (b, lq)) + # FA4 (Blackwell sm_100): preferred when available unless caller pins + # version=2/3 or env var disables. iter-5. + elif (version is None or version == 4) and _USE_FA4 and FLASH_ATTN_4_AVAILABLE: + # FA4 uses None for "no window"; FA2 used (-1, -1). + ws = ( + None if window_size[0] is None or window_size[0] < 0 else window_size[0], + None if window_size[1] is None or window_size[1] < 0 else window_size[1], + ) + out = _fa4_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + softmax_scale=softmax_scale, + causal=causal, + window_size=ws, + ) + if isinstance(out, (tuple, list)): + out = out[0] + x = out.unflatten(0, (b, lq)) + elif (version == 3 or (version is None and _USE_FA3)) and FLASH_ATTN_3_AVAILABLE: + # iter-32: FA3 (built from hopper/ source). Returns a single tensor + # at default `return_attn_probs=False`, NOT a (out, lse) tuple — the + # original `[0]` here was indexing into dim-0 of the output, giving a + # bogus (24, 128) slice. Use the return value directly. window_size + # supported by FA3 (default (-1, -1)); thread the caller's value + # through so local-attention windows work. + out = flash_attn_interface.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + seqused_q=None, + seqused_k=None, + max_seqlen_q=lq, + max_seqlen_k=lk, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + ) + if isinstance(out, (tuple, list)): + out = out[0] + x = out.unflatten(0, (b, lq)) + else: + assert FLASH_ATTN_2_AVAILABLE + x = flash_attn.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=lq, + max_seqlen_k=lk, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic).unflatten(0, (b, lq)) + + # output + return x.type(out_dtype) + + +def attention( + q, + k, + v, + q_lens=None, + k_lens=None, + dropout_p=0., + softmax_scale=None, + q_scale=None, + causal=False, + window_size=(-1, -1), + deterministic=False, + dtype=torch.bfloat16, + fa_version=None, +): + if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE: + return flash_attention( + q=q, + k=k, + v=v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + dtype=dtype, + version=fa_version, + ) + else: + if q_lens is not None or k_lens is not None: + warnings.warn( + 'Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance.' + ) + attn_mask = None + + q = q.transpose(1, 2).to(dtype) + k = k.transpose(1, 2).to(dtype) + v = v.transpose(1, 2).to(dtype) + + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p) + + out = out.transpose(1, 2).contiguous() + return out diff --git a/wan_5b/modules/causal_model.py b/wan_5b/modules/causal_model.py new file mode 100644 index 0000000..142a591 --- /dev/null +++ b/wan_5b/modules/causal_model.py @@ -0,0 +1,1882 @@ +# Adopted from https://github.com/guandeh17/Self-Forcing +# SPDX-License-Identifier: Apache-2.0 +# # Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. + +from transformers.models.x_clip.modeling_x_clip import x_clip_loss +from wan_5b.modules.attention import attention +from wan_5b.modules.model import ( + WanRMSNorm, + rope_apply, + WanLayerNorm, + WanCrossAttention, + rope_params, + sinusoidal_embedding_1d, + WanCrossAttention, + flash_attention +) +from torch.nn.attention.flex_attention import create_block_mask, flex_attention +from diffusers.configuration_utils import ConfigMixin, register_to_config +from torch.nn.attention.flex_attention import BlockMask +from diffusers.models.modeling_utils import ModelMixin +import os +import torch.nn as nn +import torch +import math +import torch.distributed as dist + +# wan 5b model compilation for flexattention +flex_attention = torch.compile( + flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs") + + +from utils.position_embedding_utils import ( + compute_temporal_freqs as _compute_temporal_freqs, + select_temporal_offset_for_sample, +) + + +# iter-21: cache freqs_i across causal_rope_apply calls within a chunk. +# All ~60 layer Q/K calls in one chunk share identical (f,h,w,start_frame, +# t_scale,temporal_offset_i,method,original_seq_len) but recompute the same +# concatenated freqs tensor each time. LRU keeps memory bounded. +# NOTE: this cache holds tensors across torch.compile step boundaries which +# is incompatible with cudagraphs (mode=reduce-overhead). If cudagraphs +# path is enabled in the future, this cache must be removed alongside +# refactoring of the KV cache scalar tensors (global_end_index, etc.). +_FREQS_I_CACHE: "dict[tuple, torch.Tensor]" = {} +_FREQS_I_CACHE_MAX = 16 +# iter-21 + iter-41: cache is on by default (iter-21 win). Set +# LLV2_FREQS_I_CACHE=0 to disable for future cudagraphs experiments (the +# cache holds tensors created inside torch.compile that get marked as +# cudagraph-pool memory; reading them on a later compile step crashes with +# "accessing tensor output of CUDAGraphs that has been overwritten"). +_FREQS_I_CACHE_ENABLED = os.environ.get("LLV2_FREQS_I_CACHE", "1") == "1" + +# iter-42: Triton fp32 RoPE kernel (utils/rope_triton.py). Default ON. +# Replaces the fp64 complex view_as_complex × complex_freqs × view_as_real +# chain with a single fused Triton kernel. Quality validated bit-exact at +# bf16 (unit test agent/rope_unit_test.py: max|Δ|=7.8e-3 = single bf16 ULP). +# Set LLV2_TRITON_ROPE=0 to revert to the fp64 path. +# When enabled, _FREQS_I_CACHE stores (freqs_i_complex, cos_f32, sin_f32); +# when disabled, stores (freqs_i_complex, None, None). +_TRITON_ROPE_ENABLED = os.environ.get("LLV2_TRITON_ROPE", "1") == "1" + +# Cudagraph experiment only. Default OFF because the out-of-place temp-KV +# construction removes mutated-input skips but is materially slower than the +# in-place temporary cache update path. +_CGRAPH_OUTPLACE_KV_ENABLED = os.environ.get("LLV2_CGRAPH_OUTPLACE_KV", "0") == "1" + +# iter-43/44: Triton fused adaLN-modulate kernel (utils/adaln_triton.py). +# Default ON after iter-44 added `@triton.autotune` over (num_warps, num_stages). +# iter-43 (no autotune) was FLAT vs iter-42 (median -1.0%, p90 +5.8%, total +# identical) — fixed config beat the eager median but jitter on tail. +# iter-44 (autotuned) is WIN: median -1.7%, p90 -1.6%, total -1.5%, FPS +1.5% +# vs iter-42, quality in run-to-run noise floor (mean|Δ|=0.68 vs noise=0.69). +# Unit test agent/adaln_unit_test.py: max|Δ|=3.1e-2 (1 bf16 ULP), mean=1.1e-3. +# Set LLV2_TRITON_ADALN=0 to fall back to eager nn.LayerNorm + Python modulate. +_TRITON_ADALN_ENABLED = os.environ.get("LLV2_TRITON_ADALN", "1") == "1" + +# iter-31: per-chunk Python-int metadata published by CausalWanModel.forward +# so attention forwards can read Python ints without `.item()` graph breaks. +# Single-thread inference assumption — overwritten before each model() call. +_CURRENT_GRID_META: "dict[str, int]" = {} + +# iter-35: removed (LOST). Consolidating duplicate .item() reads caused +# p90 latency to spike +10% — dynamo apparently traced more specialized +# paths when local vars were used in branches vs fresh .item() reads each +# time. Restored original .item() per-use pattern. + + +def causal_rope_apply(x, grid_sizes, freqs, start_frame=0, t_scale=1.0, + method="linear", original_seq_len=None, + temporal_offset=0.0): + n, c = x.size(2), x.size(3) // 2 + + # split freqs + freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) + + # loop over samples + output = [] + + # iter-47 (grad-safety fix): the Triton RoPE kernel (rope_apply_triton) is a + # raw @triton.jit op with NO autograd backward — its output is graph-detached + # (requires_grad=False). Running it under grad would silently sever gradients + # to q/k (and their LoRA). Mirror the adaLN gate (see `use_triton_adaln`): + # use Triton only when grad is OFF (inference / no_grad rollout steps); fall + # back to the differentiable fp64-complex path whenever grad is ON (training). + use_triton_rope = _TRITON_ROPE_ENABLED and not torch.is_grad_enabled() + + # iter-30: accept Python list/tuple to skip the .tolist() graph break. + # Callers that already have Python ints (sink_grid, local_grid, window_grid_sizes) + # now pass a plain list instead of `torch.tensor([[..]]).expand(...)`. + if isinstance(grid_sizes, (list, tuple)): + fwh_list = grid_sizes + else: + fwh_list = grid_sizes.tolist() + for i, (f, h, w) in enumerate(fwh_list): + seq_len = f * h * w + + # precompute multipliers — only needed for the fp64 complex path. + # iter-42: skip the bf16→fp64 cast + view_as_complex when the Triton + # kernel will be used (it consumes bf16 directly). + # iter-47: gate on use_triton_rope (not the raw flag) so the complex x_i + # IS precomputed whenever we fall back to the differentiable path (training). + if not use_triton_rope: + x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape( + seq_len, n, -1, 2)) + temporal_offset_i = select_temporal_offset_for_sample( + temporal_offset, i, f, start_frame=start_frame) + + # iter-21: cache freqs_i. iter-41: gate the cache behind + # LLV2_FREQS_I_CACHE=1 (default off). The cache stores tensors + # created inside torch.compile, which cudagraph allocator considers + # owned by the per-step memory pool — reading them on a later step + # races with the pool's reuse. Disabling the cache unblocks + # `mode=reduce-overhead` for cudagraphs; the recomputation cost is + # tiny (60 layer calls × per-chunk concat ≈ 0.5% wall) compared to + # the cudagraphs unlock potential. + if _FREQS_I_CACHE_ENABLED: + if torch.is_tensor(temporal_offset_i): + if temporal_offset_i.ndim == 0: + offset_key = float(temporal_offset_i.item()) + else: + offset_key = ("tensor", id(temporal_offset_i)) + else: + offset_key = float(temporal_offset_i) + cache_key = ( + f, h, w, start_frame, t_scale, method, + original_seq_len, offset_key, x.device.type, x.device.index, + use_triton_rope, # iter-47: separate cached repr for grad/no-grad + ) + cache_entry = _FREQS_I_CACHE.get(cache_key) + else: + cache_entry = None + cache_key = None + + if cache_entry is None: + temporal_freqs = _compute_temporal_freqs( + freqs[0], f, start_frame, t_scale, x.device, + method=method, original_seq_len=original_seq_len, + temporal_offset=temporal_offset_i) + freqs_i_complex = torch.cat([ + temporal_freqs.view(f, 1, 1, -1).expand(f, h, w, -1), + freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1), + ], dim=-1).reshape(seq_len, 1, -1) + if use_triton_rope: + # iter-42: store (cos, sin) fp32 derived once; freqs_i_complex + # itself only kept for the legacy fp64 path. + from utils.rope_triton import _split_complex_to_cos_sin + cos_f32, sin_f32 = _split_complex_to_cos_sin(freqs_i_complex) + cache_entry = (freqs_i_complex, cos_f32, sin_f32) + else: + cache_entry = (freqs_i_complex, None, None) + if _FREQS_I_CACHE_ENABLED: + if len(_FREQS_I_CACHE) >= _FREQS_I_CACHE_MAX: + _FREQS_I_CACHE.pop(next(iter(_FREQS_I_CACHE))) + _FREQS_I_CACHE[cache_key] = cache_entry + freqs_i, cos_f32, sin_f32 = cache_entry + + # apply rotary embedding + if use_triton_rope: + # iter-42: Triton fp32 kernel — replaces the fp64 complex128 path. + # iter-46: kernel takes full x[i] + seq_len and emits rotated-or- + # passthrough output in a single launch, eliminating the + # `.contiguous()` slice + outer `torch.cat`. Bit-exact preserved. + from utils.rope_triton import rope_apply_triton + x_i = rope_apply_triton(x[i], cos_f32, sin_f32, seq_len=seq_len) + else: + x_i = torch.view_as_real(x_i * freqs_i).flatten(2) + x_i = torch.cat([x_i, x[i, seq_len:]]) + + # append to collection + output.append(x_i) + return torch.stack(output).type_as(x) + + +class MultiShotT2VCrossAttention(WanCrossAttention): + + def forward(self, x, context, context_lens, is_teacher_forcing=False, crossattn_cache=None): + r""" + Args: + x(Tensor): Shape [B, L1, C] + context(Tensor): Shape [B * num_chunks, L2, C] + context_lens(Tensor): Shape [B] or [B * num_chunks] + crossattn_cache (List[dict], *optional*): Contains the cached key and value tensors for context embedding. + """ + # Original batch size (videos) + b_orig, L1, C = x.size() + n, d = self.num_heads, self.head_dim + + # Effective batch size for cross-attention (videos * chunks) + b_ctx = context.size(0) + assert b_ctx % b_orig == 0, f"context batch ({b_ctx}) must be a multiple of x batch ({b_orig})" + num_chunks = b_ctx // b_orig + + # Prepare context_lens for [B * num_chunks] if needed + if context_lens is not None and context_lens.numel() == b_orig: + context_lens = context_lens.repeat_interleave(num_chunks) + elif context_lens is not None: + assert context_lens.numel() == b_ctx, \ + f"context_lens must have length {b_orig} or {b_ctx}, got {context_lens.numel()}" + # Helper to run standard cross-attention on a given x_chunk of shape [B * num_chunks, L_chunk, C] + def _cross_attend(x_chunk): + b_eff, L_chunk, _ = x_chunk.size() + + # compute query, key, value + q = self.norm_q(self.q(x_chunk)).view(b_eff, -1, n, d) + + # iter-24: Bypass crossattn_cache. Cached K/V tensors escape the + # cudagraph memory pool across torch.compile step boundaries and + # block `mode=reduce-overhead`. Per-call recompute cost is tiny + # (~1.7us / call in NVFP4 × ~11.5k calls/prompt ≈ 19 ms total), + # for cudagraphs unlock of the 28% wall-time gap. + k = self.norm_k(self.k(context)).view(b_eff, -1, n, d) + v = self.v(context).view(b_eff, -1, n, d) + + # compute attention + x_attn = flash_attention(q, k, v, k_lens=context_lens) + + # output projection + x_attn = x_attn.flatten(2) + x_attn = self.o(x_attn) + return x_attn + + if not is_teacher_forcing: + # ------------------------------- + # Regular multi-shot: all tokens attend text, we just chunk along L1 + # x: [B, L1, C] -> [B * num_chunks, L1 / num_chunks, C] + # ------------------------------- + assert L1 % num_chunks == 0, \ + f"L1 ({L1}) must be divisible by num_chunks ({num_chunks})" + tokens_per_chunk = L1 // num_chunks + + x_chunked = x.view(b_orig, num_chunks, tokens_per_chunk, C) + x_chunked = x_chunked.reshape(b_ctx, tokens_per_chunk, C) + + x_attn = _cross_attend(x_chunked) # [B * num_chunks, tokens_per_chunk, C] + + # reshape back to [B, L1, C] + x_attn = x_attn.view(b_orig, num_chunks, tokens_per_chunk, C) + x_attn = x_attn.reshape(b_orig, L1, C) + return x_attn + + # ------------------------------- + # Teacher forcing: + # x is typically [B, 2 * L_tf, C], where the first half is clean and + # the second half is noisy. Apply multi-shot cross-attention to both + # halves. + # ------------------------------- + assert L1 % 2 == 0, f"In teacher-forcing mode, L1 ({L1}) should be even." + half = L1 // 2 + x_clean = x[:, :half, :] # [B, L_tf, C] + x_noisy = x[:, half:, :] # [B, L_tf, C] + + def _chunk_and_attend(x_part): + L_part = x_part.size(1) + assert L_part % num_chunks == 0, \ + f"Segment length ({L_part}) must be divisible by num_chunks ({num_chunks})" + tokens_per_chunk = L_part // num_chunks + + # [B, L_part, C] -> [B * num_chunks, L_part / num_chunks, C] + x_chunked = x_part.view(b_orig, num_chunks, tokens_per_chunk, C) + x_chunked = x_chunked.reshape(b_ctx, tokens_per_chunk, C) + + x_attn = _cross_attend(x_chunked) # [B * num_chunks, tokens_per_chunk, C] + x_attn = x_attn.view(b_orig, num_chunks, tokens_per_chunk, C) + x_attn = x_attn.reshape(b_orig, L_part, C) + return x_attn + + x_clean_attn = _chunk_and_attend(x_clean) + x_noisy_attn = _chunk_and_attend(x_noisy) + + # Reassemble the full sequence from cross-attended clean and noisy halves. + x_out = torch.cat([x_clean_attn, x_noisy_attn], dim=1) # [B, L1, C] + return x_out + + +class CausalWanSelfAttention(nn.Module): + + def __init__(self, + dim, + num_heads, + local_attn_size=-1, + sink_size=0, + qk_norm=True, + eps=1e-6): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.local_attn_size = local_attn_size if local_attn_size != -1 else 24 + self.sink_size = sink_size + self.global_sink_size = 0 + self.qk_norm = qk_norm + self.eps = eps + self.max_attention_size = 24 * 880 if local_attn_size == -1 else local_attn_size * 880 + + # layers + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + + def forward( + self, + x, + seq_lens, + grid_sizes, + freqs, + block_mask, + kv_cache=None, + current_start=0, + cache_start=None, + t_scale=1.0, + use_relative_rope=False, + method="linear", + original_seq_len=None, + temporal_offset=0.0, + ): + r""" + Args: + x(Tensor): Shape [B, L, num_heads, C / num_heads] + seq_lens(Tensor): Shape [B] + grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + block_mask (BlockMask) + t_scale (float): Temporal RoPE interpolation scale. <1.0 compresses positions. + use_relative_rope (bool): If True, store raw K in cache and apply RoPE + with window-relative positions at attention time. + method (str): RoPE method. This release supports "linear". + original_seq_len (int): Unused by the release linear RoPE path. + temporal_offset (float): Multi-shot RoPE offset (shot_index * phi). + """ + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + if cache_start is None: + cache_start = current_start + + # query, key, value function + def qkv_fn(x): + q = self.norm_q(self.q(x)).view(b, s, n, d) + k = self.norm_k(self.k(x)).view(b, s, n, d) + v = self.v(x).view(b, s, n, d) + return q, k, v + + q, k, v = qkv_fn(x) + + if kv_cache is None: + # Teacher-forcing training doubles sequence length with clean/noisy halves. + is_tf = (s == seq_lens[0].item() * 2) + if is_tf: + q_chunk = torch.chunk(q, 2, dim=1) + k_chunk = torch.chunk(k, 2, dim=1) + roped_query = [] + roped_key = [] + # rope should be same for clean and noisy parts + for ii in range(2): + rq = rope_apply(q_chunk[ii], grid_sizes, freqs, t_scale=t_scale, + method=method, original_seq_len=original_seq_len, + temporal_offset=temporal_offset).type_as(v) + rk = rope_apply(k_chunk[ii], grid_sizes, freqs, t_scale=t_scale, + method=method, original_seq_len=original_seq_len, + temporal_offset=temporal_offset).type_as(v) + roped_query.append(rq) + roped_key.append(rk) + + roped_query = torch.cat(roped_query, dim=1) + roped_key = torch.cat(roped_key, dim=1) + + padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1] + padded_roped_query = torch.cat( + [roped_query, + torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]], + device=q.device, dtype=v.dtype)], + dim=1 + ) + + padded_roped_key = torch.cat( + [roped_key, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]], + device=k.device, dtype=v.dtype)], + dim=1 + ) + + padded_v = torch.cat( + [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]], + device=v.device, dtype=v.dtype)], + dim=1 + ) + + x = flex_attention( + query=padded_roped_query.transpose(2, 1), + key=padded_roped_key.transpose(2, 1), + value=padded_v.transpose(2, 1), + block_mask=block_mask + ) + x = x[:, :, :(-padded_length)] if padded_length > 0 else x + x = x.transpose(2, 1) + + else: + roped_query = rope_apply(q, grid_sizes, freqs, t_scale=t_scale, + method=method, original_seq_len=original_seq_len, + temporal_offset=temporal_offset).type_as(v) + roped_key = rope_apply(k, grid_sizes, freqs, t_scale=t_scale, + method=method, original_seq_len=original_seq_len, + temporal_offset=temporal_offset).type_as(v) + + padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1] + padded_roped_query = torch.cat( + [roped_query, + torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]], + device=q.device, dtype=v.dtype)], + dim=1 + ) + + padded_roped_key = torch.cat( + [roped_key, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]], + device=k.device, dtype=v.dtype)], + dim=1 + ) + + padded_v = torch.cat( + [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]], + device=v.device, dtype=v.dtype)], + dim=1 + ) + + x = flex_attention( + query=padded_roped_query.transpose(2, 1), + key=padded_roped_key.transpose(2, 1), + value=padded_v.transpose(2, 1), + block_mask=block_mask + ) + x = x[:, :, :(-padded_length)] if padded_length > 0 else x + x = x.transpose(2, 1) + else: + # iter-31: read Python ints from module-level dict (set by + # CausalWanModel.forward) instead of `.item()` calls on + # grid_sizes, removing 4 graph breaks per attention forward. + if _CURRENT_GRID_META: + frame_seqlen = _CURRENT_GRID_META["frame_seqlen"] + num_new_frames = _CURRENT_GRID_META["num_new_frames"] + h = _CURRENT_GRID_META["h"] + w = _CURRENT_GRID_META["w"] + else: + frame_seqlen = math.prod(grid_sizes[0][1:]).item() + num_new_frames = grid_sizes[0][0].item() + h, w = grid_sizes[0][1].item(), grid_sizes[0][2].item() + num_new_tokens = q.shape[1] + current_end = current_start + num_new_tokens + # iter-30: build Python-int grid once; pass to all rope_apply calls + # below so they skip the .tolist() graph break. + b = q.shape[0] + grid_py = [(num_new_frames, h, w)] * b + + if not use_relative_rope: + current_start_frame = current_start // frame_seqlen + roped_query = causal_rope_apply( + q, grid_py, freqs, start_frame=current_start_frame, t_scale=t_scale, + method=method, original_seq_len=original_seq_len, + temporal_offset=temporal_offset).type_as(v) + roped_key = causal_rope_apply( + k, grid_py, freqs, start_frame=current_start_frame, t_scale=t_scale, + method=method, original_seq_len=original_seq_len, + temporal_offset=temporal_offset).type_as(v) + key_to_cache = roped_key + else: + key_to_cache = k + + sink_tokens = self.sink_size * frame_seqlen + global_sink_tokens = getattr(self, "global_sink_size", 0) * frame_seqlen + is_quantized_cache = kv_cache.get("quantized", False) + if is_quantized_cache: + kv_cache_size = kv_cache["max_blocks"] * kv_cache["block_token_size"] + else: + kv_cache_size = kv_cache["k"].shape[1] + + # ----- global + multi-shot pinned-sink support ----- + # Two protection mechanisms (independent, both optional): + # * global_sink_tokens: first N frames are permanently anchored + # (set via global_sink_size; never moves, always attended). + # * pinned region (pinned_start/pinned_len): multi-shot sink put + # on a scene cut. The pinned chunk lives at its original buffer + # position; rolling shifts non-pinned data around it. + # effective_sink = leading buffer prefix that rolling MUST keep: + # pinned right after global (pinned_start == global_sink_tokens) + # -> global_sink_tokens + pinned_len + # pinned elsewhere (floating) + # -> global_sink_tokens + # no pinned + # -> max(global_sink_tokens, sink_tokens) # legacy compat + # iter-39: read pinned state from _CURRENT_GRID_META (published + # once per chunk in CausalWanModel._forward_inference). Falls + # back to `.item()` if the dict was not initialized (e.g. unit + # test exercising the attention block directly). + if _CURRENT_GRID_META and "pinned_start" in _CURRENT_GRID_META: + pinned_start_val = _CURRENT_GRID_META["pinned_start"] + pinned_len_val = _CURRENT_GRID_META["pinned_len"] + else: + pinned_start_t = kv_cache.get("pinned_start", None) + pinned_len_val = 0 + if pinned_start_t is not None and hasattr(pinned_start_t, 'item'): + pinned_start_val = pinned_start_t.item() + pinned_len_val = kv_cache["pinned_len"].item() + else: + pinned_start_val = -1 + has_pinned = pinned_start_val >= 0 and pinned_len_val > 0 + if has_pinned and pinned_start_val == global_sink_tokens: + effective_sink = global_sink_tokens + pinned_len_val + elif has_pinned: + effective_sink = global_sink_tokens + else: + effective_sink = max(global_sink_tokens, sink_tokens) + + # iter-39: read cache indices from _CURRENT_GRID_META (published + # by CausalWanModel._forward_inference) to avoid 6+ `.item()` + # syncs per block forward. Falls back to .item() when the dict + # is not initialized (direct attention-block unit tests). + if _CURRENT_GRID_META and "global_end_index" in _CURRENT_GRID_META: + _cache_global_end = _CURRENT_GRID_META["global_end_index"] + _cache_local_end = _CURRENT_GRID_META["local_end_index"] + else: + _cache_global_end = kv_cache["global_end_index"].item() + _cache_local_end = kv_cache["local_end_index"].item() + + cache_update_info = None + if self.local_attn_size != -1 and (current_end > _cache_global_end) and ( + num_new_tokens + _cache_local_end > kv_cache_size): + num_evicted_tokens = num_new_tokens + _cache_local_end - kv_cache_size + num_rolled_tokens = _cache_local_end - num_evicted_tokens - effective_sink + + local_end_index = _cache_local_end + current_end - \ + _cache_global_end - num_evicted_tokens + local_start_index = local_end_index - num_new_tokens + + if is_quantized_cache: + from utils.quant import dequantize_kv_cache, k_smooth + + max_blks = int(kv_cache["max_blocks"]) + blk_sz = int(kv_cache["block_token_size"]) + cache_k = dequantize_kv_cache( + kv_cache["k"], max_blks, self.num_heads, blk_sz, v.dtype, v.device + ) + cache_v = dequantize_kv_cache( + kv_cache["v"], max_blks, self.num_heads, blk_sz, v.dtype, v.device + ) + new_k_for_cache = k_smooth(key_to_cache) + else: + cache_k = kv_cache["k"] + cache_v = kv_cache["v"] + new_k_for_cache = key_to_cache + + if _CGRAPH_OUTPLACE_KV_ENABLED: + # Cudagraph experiment: build the post-roll cache view + # out-of-place. Slice assignment here forces Inductor + # cudagraph partitions to mutate inputs. + temp_k = torch.cat([ + cache_k[:, :effective_sink], + cache_k[:, effective_sink + num_evicted_tokens: + effective_sink + num_evicted_tokens + num_rolled_tokens], + new_k_for_cache, + ], dim=1) + temp_v = torch.cat([ + cache_v[:, :effective_sink], + cache_v[:, effective_sink + num_evicted_tokens: + effective_sink + num_evicted_tokens + num_rolled_tokens], + v, + ], dim=1) + else: + temp_k = cache_k if is_quantized_cache else cache_k.clone() + temp_v = cache_v if is_quantized_cache else cache_v.clone() + + temp_k[:, effective_sink:effective_sink + num_rolled_tokens] = \ + temp_k[:, effective_sink + num_evicted_tokens:effective_sink + num_evicted_tokens + num_rolled_tokens].clone() + temp_v[:, effective_sink:effective_sink + num_rolled_tokens] = \ + temp_v[:, effective_sink + num_evicted_tokens:effective_sink + num_evicted_tokens + num_rolled_tokens].clone() + + temp_k[:, local_start_index:local_end_index] = new_k_for_cache + temp_v[:, local_start_index:local_end_index] = v + + # When pinned is "floating" (lives outside effective_sink), the + # rolling shifted non-pinned data left by num_evicted_tokens; + # the pinned anchor must follow that shift to keep tracking the + # same data. When pinned sits inside effective_sink (i.e. right + # after the global region), it is part of the protected prefix + # and rolling does not move it. + pinned_shift = num_evicted_tokens if (has_pinned and pinned_start_val >= effective_sink) else 0 + + cache_update_info = { + "action": "roll_and_insert", + "sink_tokens": effective_sink, + "num_rolled_tokens": num_rolled_tokens, + "num_evicted_tokens": num_evicted_tokens, + "local_start_index": local_start_index, + "local_end_index": local_end_index, + "new_k": key_to_cache, + "new_v": v, + "current_end": current_end, + "pinned_shift": pinned_shift, + } + + else: + # iter-39: reuse the dict-cached scalars from above. + local_end_index = _cache_local_end + current_end - _cache_global_end + local_start_index = local_end_index - num_new_tokens + + if is_quantized_cache: + from utils.quant import dequantize_kv_cache, k_smooth + + new_k_for_cache = k_smooth(key_to_cache) + if local_start_index == 0: + temp_k = new_k_for_cache + temp_v = v + else: + max_blks = int(kv_cache["max_blocks"]) + blk_sz = int(kv_cache["block_token_size"]) + cache_k = dequantize_kv_cache( + kv_cache["k"], max_blks, self.num_heads, blk_sz, v.dtype, v.device + ) + cache_v = dequantize_kv_cache( + kv_cache["v"], max_blks, self.num_heads, blk_sz, v.dtype, v.device + ) + if _CGRAPH_OUTPLACE_KV_ENABLED: + temp_k = torch.cat([cache_k[:, :local_start_index], new_k_for_cache], dim=1) + temp_v = torch.cat([cache_v[:, :local_start_index], v], dim=1) + else: + temp_k = cache_k + temp_v = cache_v + if not _CGRAPH_OUTPLACE_KV_ENABLED: + temp_k[:, local_start_index:local_end_index] = new_k_for_cache + temp_v[:, local_start_index:local_end_index] = v + else: + if _CGRAPH_OUTPLACE_KV_ENABLED: + temp_k = torch.cat([kv_cache["k"][:, :local_start_index], key_to_cache], dim=1) + temp_v = torch.cat([kv_cache["v"][:, :local_start_index], v], dim=1) + else: + temp_k = kv_cache["k"].clone() + temp_v = kv_cache["v"].clone() + temp_k[:, local_start_index:local_end_index] = key_to_cache + temp_v[:, local_start_index:local_end_index] = v + + cache_update_info = { + "action": "direct_insert", + "local_start_index": local_start_index, + "local_end_index": local_end_index, + "new_k": key_to_cache, + "new_v": v, + "current_end": current_end, + "pinned_shift": 0, + } + + window_start = max(0, local_end_index - self.max_attention_size) + + # Build the K/V actually attended over. + # Cases: + # (a) prepend_sink : effective_sink > 0 and out of window + # -> prepend [:effective_sink] (covers global + # and any pinned-merged-to-front) + # (b) prepend_pinned: a floating pinned region (pinned_start + # >= effective_sink) lives outside the window + # -> additionally prepend that pinned slice + # (c) otherwise : plain sliding window + # Note (a) and (b) are not mutually exclusive: when global is + # enabled AND there is a separate floating pinned region outside + # the window, both prefixes must be prepended. + prepend_sink = effective_sink > 0 and window_start > 0 + prepend_pinned = ( + has_pinned and pinned_start_val >= effective_sink + and pinned_start_val < window_start + ) + + if prepend_sink and prepend_pinned: + # [global+sink] + [pinned] + [local window] + extra = effective_sink + pinned_len_val + effective_local_size = self.max_attention_size - extra + local_window_start = max(effective_sink, local_end_index - effective_local_size) + window_k = torch.cat([ + temp_k[:, :effective_sink], + temp_k[:, pinned_start_val:pinned_start_val + pinned_len_val], + temp_k[:, local_window_start:local_end_index], + ], dim=1) + window_v = torch.cat([ + temp_v[:, :effective_sink], + temp_v[:, pinned_start_val:pinned_start_val + pinned_len_val], + temp_v[:, local_window_start:local_end_index], + ], dim=1) + elif prepend_sink: + effective_local_size = self.max_attention_size - effective_sink + local_window_start = max(effective_sink, local_end_index - effective_local_size) + window_k = torch.cat([temp_k[:, :effective_sink], temp_k[:, local_window_start:local_end_index]], dim=1) + window_v = torch.cat([temp_v[:, :effective_sink], temp_v[:, local_window_start:local_end_index]], dim=1) + elif prepend_pinned: + effective_local_size = self.max_attention_size - pinned_len_val + local_window_start = max(0, local_end_index - effective_local_size) + window_k = torch.cat( + [temp_k[:, pinned_start_val:pinned_start_val + pinned_len_val], + temp_k[:, local_window_start:local_end_index]], dim=1) + window_v = torch.cat( + [temp_v[:, pinned_start_val:pinned_start_val + pinned_len_val], + temp_v[:, local_window_start:local_end_index]], dim=1) + else: + window_k = temp_k[:, window_start:local_end_index] + window_v = temp_v[:, window_start:local_end_index] + + if use_relative_rope: + if prepend_sink: + # Sink and local window tokens get separate RoPE in a + # virtual contiguous layout: [sink_frames | local_frames]. + sink_frame_count = effective_sink // frame_seqlen + local_tokens = window_k.shape[1] - effective_sink + local_frame_count = local_tokens // frame_seqlen + combined_frames = sink_frame_count + local_frame_count + + # iter-30: pass Python list instead of expanded tensor; + # causal_rope_apply skips .tolist() graph break this way. + sink_grid = [(sink_frame_count, h, w)] * b + roped_sink_k = causal_rope_apply( + window_k[:, :effective_sink], sink_grid, freqs, + start_frame=0, t_scale=t_scale, + method=method, original_seq_len=original_seq_len, + ).type_as(v) + + local_grid = [(local_frame_count, h, w)] * b + roped_local_k = causal_rope_apply( + window_k[:, effective_sink:], local_grid, freqs, + start_frame=sink_frame_count, t_scale=t_scale, + method=method, original_seq_len=original_seq_len, + ).type_as(v) + + roped_window_k = torch.cat([roped_sink_k, roped_local_k], dim=1) + + q_start_frame = combined_frames - num_new_frames + roped_query = causal_rope_apply( + q, grid_py, freqs, + start_frame=q_start_frame, t_scale=t_scale, + method=method, original_seq_len=original_seq_len, + ).type_as(v) + else: + window_tokens = window_k.shape[1] + window_frames = window_tokens // frame_seqlen + + # iter-30: Python list to skip .tolist() break. + window_grid_sizes = [(window_frames, h, w)] * b + + roped_window_k = causal_rope_apply( + window_k, window_grid_sizes, freqs, + start_frame=0, t_scale=t_scale, + method=method, original_seq_len=original_seq_len, + ).type_as(v) + + q_start_frame = window_frames - num_new_frames + roped_query = causal_rope_apply( + q, grid_py, freqs, + start_frame=q_start_frame, t_scale=t_scale, + method=method, original_seq_len=original_seq_len, + ).type_as(v) + + x = attention(roped_query, roped_window_k, window_v) + else: + x = attention(roped_query, window_k, window_v) + + # output + x = x.flatten(2) + x = self.o(x) + + # Return both output and cache update info + if kv_cache is not None: + return x, (current_end, local_end_index, cache_update_info) + else: + return x + + +class CausalWanAttentionBlock(nn.Module): + + def __init__(self, + dim, + ffn_dim, + num_heads, + local_attn_size=-1, + sink_size=0, + qk_norm=True, + cross_attn_norm=False, + eps=1e-6): + super().__init__() + self.dim = dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + self.local_attn_size = local_attn_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + + # layers + self.norm1 = WanLayerNorm(dim, eps) + self.self_attn = CausalWanSelfAttention(dim, num_heads, local_attn_size, sink_size, qk_norm, eps) + self.norm3 = WanLayerNorm( + dim, eps, + elementwise_affine=True) if cross_attn_norm else nn.Identity() + self.cross_attn = MultiShotT2VCrossAttention(dim, num_heads, (-1, -1), qk_norm, eps) + self.norm2 = WanLayerNorm(dim, eps) + self.ffn = nn.Sequential( + nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'), + nn.Linear(ffn_dim, dim)) + + # modulation + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + x, + e, + seq_lens, + grid_sizes, + freqs, + context, + context_lens, + block_mask, + kv_cache=None, + crossattn_cache=None, + current_start=0, + cache_start=None, + t_scale=1.0, + use_relative_rope=False, + method="linear", + original_seq_len=None, + temporal_offset=0.0, + ): + r""" + Args: + x(Tensor): Shape [B, L, C] + e(Tensor): Shape [B, F, 6, C] + seq_lens(Tensor): Shape [B], length of each sequence in batch + grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + t_scale (float): Temporal RoPE interpolation scale. <1.0 compresses positions. + use_relative_rope (bool): If True, use window-relative RoPE positions in KV cache path. + method (str): RoPE method. This release supports "linear". + original_seq_len (int): Unused by the release linear RoPE path. + temporal_offset (float): Multi-shot RoPE offset (shot_index * phi). + """ + num_frames, frame_seqlen = e.shape[1], x.shape[1] // e.shape[1] + e = (self.modulation.unsqueeze(1) + e).chunk(6, dim=2) + use_triton_adaln = _TRITON_ADALN_ENABLED and not torch.is_grad_enabled() + + # self-attention + if use_triton_adaln: + # iter-43: fused LayerNorm + (1+e[1])*x + e[0] in one Triton kernel. + from utils.adaln_triton import adaln_modulate_triton + modulated_x = adaln_modulate_triton( + x, e[1], e[0], frame_seqlen, + eps=self.norm1.eps, add_one_to_scale=True, + ) + else: + modulated_x = (self.norm1(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1 + e[1]) + e[0]).flatten(1, 2) + self_attn_result = self.self_attn( + modulated_x, + seq_lens, grid_sizes, + freqs, block_mask, kv_cache, current_start, cache_start, t_scale=t_scale, + use_relative_rope=use_relative_rope, + method=method, original_seq_len=original_seq_len, + temporal_offset=temporal_offset) + + if kv_cache is not None: + y, cache_update_info = self_attn_result + else: + y = self_attn_result + cache_update_info = None + x = x + (y.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * e[2]).flatten(1, 2) + + # cross-attention & ffn function + # iter-40: avoid `seq_lens[0].item()` graph break. seq_lens[0] equals + # num_new_frames * frame_seqlen at inference time, both of which are + # Python ints already in _CURRENT_GRID_META (published by iter-31). + # `is_tf` is True only in teacher-forcing training where x.shape[1] + # is the doubled (clean+noisy) sequence — never at inference. + if _CURRENT_GRID_META and "frame_seqlen" in _CURRENT_GRID_META: + seq_len_py = ( + _CURRENT_GRID_META["frame_seqlen"] + * _CURRENT_GRID_META["num_new_frames"] + ) + is_tf = (x.shape[1] == seq_len_py * 2) + else: + is_tf = (x.shape[1] == seq_lens[0].item() * 2) + def cross_attn_ffn(x, context, context_lens, e, crossattn_cache=None): + x = x + self.cross_attn(self.norm3(x), context, + context_lens, crossattn_cache=crossattn_cache, is_teacher_forcing=is_tf) + if use_triton_adaln: + # iter-43: fused LayerNorm + (1+e[4])*x + e[3] in one Triton kernel. + from utils.adaln_triton import adaln_modulate_triton + ffn_in = adaln_modulate_triton( + x, e[4], e[3], frame_seqlen, + eps=self.norm2.eps, add_one_to_scale=True, + ) + else: + ffn_in = (self.norm2(x).unflatten(dim=1, sizes=(num_frames, + frame_seqlen)) * (1 + e[4]) + e[3]).flatten(1, 2) + y = self.ffn(ffn_in) + x = x + (y.unflatten(dim=1, sizes=(num_frames, + frame_seqlen)) * e[5]).flatten(1, 2) + return x + + x = cross_attn_ffn(x, context, context_lens, e, crossattn_cache) + + if cache_update_info is not None: + # cache_update_info is already formatted as + # (current_end, local_end_index, cache_update_info). + return x, cache_update_info + else: + return x + + +class CausalHead(nn.Module): + + def __init__(self, dim, out_dim, patch_size, eps=1e-6): + super().__init__() + self.dim = dim + self.out_dim = out_dim + self.patch_size = patch_size + self.eps = eps + + # layers + out_dim = math.prod(patch_size) * out_dim + self.norm = WanLayerNorm(dim, eps) + self.head = nn.Linear(dim, out_dim) + + # modulation + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x, e): + r""" + Args: + x(Tensor): Shape [B, L1, C] + e(Tensor): Shape [B, F, 1, C] + """ + num_frames, frame_seqlen = e.shape[1], x.shape[1] // e.shape[1] + e = (self.modulation.unsqueeze(1) + e).chunk(2, dim=2) + x = (self.head(self.norm(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1 + e[1]) + e[0])) + return x + + +class CausalWanModel(ModelMixin, ConfigMixin): + r""" + Wan diffusion backbone supporting both text-to-video and image-to-video with causal attention. + """ + + ignore_for_config = [ + 'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim', 'window_size' + ] + _no_split_modules = ['WanAttentionBlock'] + _supports_gradient_checkpointing = True + + @register_to_config + def __init__(self, + model_type='t2v', + patch_size=(1, 2, 2), + text_len=512, + in_dim=16, + dim=2048, + ffn_dim=8192, + freq_dim=256, + text_dim=4096, + out_dim=16, + num_heads=16, + num_layers=32, + local_attn_size=-1, + sink_size=0, + num_frame_per_block=1, + qk_norm=True, + cross_attn_norm=True, + eps=1e-6): + r""" + Initialize the diffusion model backbone. + + Args: + model_type (`str`, *optional*, defaults to 't2v'): + Model variant - 't2v' (text-to-video), 'i2v' (image-to-video), or 'ti2v' + patch_size (`tuple`, *optional*, defaults to (1, 2, 2)): + 3D patch dimensions for video embedding (t_patch, h_patch, w_patch) + text_len (`int`, *optional*, defaults to 512): + Fixed length for text embeddings + in_dim (`int`, *optional*, defaults to 16): + Input video channels (C_in) + dim (`int`, *optional*, defaults to 2048): + Hidden dimension of the transformer + ffn_dim (`int`, *optional*, defaults to 8192): + Intermediate dimension in feed-forward network + freq_dim (`int`, *optional*, defaults to 256): + Dimension for sinusoidal time embeddings + text_dim (`int`, *optional*, defaults to 4096): + Input dimension for text embeddings + out_dim (`int`, *optional*, defaults to 16): + Output video channels (C_out) + num_heads (`int`, *optional*, defaults to 16): + Number of attention heads + num_layers (`int`, *optional*, defaults to 32): + Number of transformer blocks + local_attn_size (`int`, *optional*, defaults to -1): + Window size for temporal local attention (-1 indicates global attention) + sink_size (`int`, *optional*, defaults to 0): + Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache + qk_norm (`bool`, *optional*, defaults to True): + Enable query/key normalization + cross_attn_norm (`bool`, *optional*, defaults to False): + Enable cross-attention normalization + eps (`float`, *optional*, defaults to 1e-6): + Epsilon value for normalization layers + """ + + super().__init__() + + assert model_type in ['t2v', 'i2v', 'ti2v'] + self.model_type = model_type + + self.patch_size = patch_size + self.text_len = text_len + self.in_dim = in_dim + self.dim = dim + self.ffn_dim = ffn_dim + self.freq_dim = freq_dim + self.text_dim = text_dim + self.out_dim = out_dim + self.num_heads = num_heads + self.num_layers = num_layers + self.local_attn_size = local_attn_size + self.sink_size = sink_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + + # embeddings + self.patch_embedding = nn.Conv3d( + in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'), + nn.Linear(dim, dim)) + + self.time_embedding = nn.Sequential( + nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + self.time_projection = nn.Sequential( + nn.SiLU(), nn.Linear(dim, dim * 6)) + + # blocks + self.blocks = nn.ModuleList([ + CausalWanAttentionBlock(dim, ffn_dim, num_heads, + local_attn_size, sink_size, qk_norm, cross_attn_norm, eps) + for _ in range(num_layers) + ]) + + # head + self.head = CausalHead(dim, out_dim, patch_size, eps) + + # buffers (don't use register_buffer otherwise dtype will be changed in to()) + assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0 + d = dim // num_heads + self.freqs = torch.cat([ + rope_params(1024, d - 4 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + rope_params(1024, 2 * (d // 6)) + ], + dim=1) + + # initialize weights + self.init_weights() + + self.gradient_checkpointing = False + + self.block_mask = None + self._block_mask_batch_size = 0 + + self.num_frame_per_block = num_frame_per_block + self.independent_first_frame = False + self.t_scale = 1.0 + self.use_relative_rope = False + self.rope_method = "linear" + self.original_seq_len = None + self.rope_temporal_offset = 0.0 + self.kv_quant_config = None + + def _set_gradient_checkpointing(self, module, value=False): + self.gradient_checkpointing = value + + @staticmethod + def _prepare_blockwise_causal_attn_mask_i2v( + device: torch.device | str, num_frames: int = 31, + frame_seqlen: int = 880, num_frame_per_block=3, + batch_size=None, + ) -> BlockMask: + + """ + we will divide the token sequence into the following format + [1 latent frame] [N latent frame] ... [N latent frame] + The first frame is separated out to support I2V generation + We use flexattention to construct the attention mask + """ + total_length = num_frames * frame_seqlen + + # we do right padding to get to a multiple of 128 + padded_length = math.ceil(total_length / 128) * 128 - total_length + + ends = torch.zeros(total_length + padded_length, + device=device, dtype=torch.long) + + # special handling for the first frame + ends[:frame_seqlen] = frame_seqlen + + # Block-wise causal mask will attend to all elements that are before the end of the current chunk + frame_indices = torch.arange( + start=frame_seqlen, + end=total_length, + step=frame_seqlen * num_frame_per_block, + device=device + ) + + for idx, tmp in enumerate(frame_indices): + ends[tmp:tmp + frame_seqlen * num_frame_per_block] = tmp + \ + frame_seqlen * num_frame_per_block + + def attention_mask(b, h, q_idx, kv_idx): + is_real_q = q_idx < total_length + is_real_k = kv_idx < total_length + return (q_idx == kv_idx) | (is_real_q & is_real_k & (kv_idx < ends[q_idx])) + + block_mask = create_block_mask(attention_mask, B=batch_size, H=None, Q_LEN=total_length + padded_length, + KV_LEN=total_length + padded_length, _compile=False, device=device) + return block_mask + + @staticmethod + def _prepare_blockwise_causal_attn_mask( + device: torch.device | str, num_frames: int = 31, + frame_seqlen: int = 880, num_frame_per_block=1, + batch_size=None, + ) -> BlockMask: + """ + Block-wise causal mask. The mask is defined only by the AR chunk size: + a token can attend to all tokens before the end of its current + num_frame_per_block chunk. + """ + print(f"num_frame_per_block: {num_frame_per_block}") + total_length = num_frames * frame_seqlen + padded_length = math.ceil(total_length / 128) * 128 - total_length + + block_size = frame_seqlen * num_frame_per_block + + def attention_mask(b, h, q_idx, kv_idx): + # Apply only to real tokens in [0, total_length); padding keeps only + # self-loops. + is_real_q = q_idx < total_length + is_real_k = kv_idx < total_length + + # End position of the block containing the current token. + current_block_end = ((q_idx // block_size) + 1) * block_size + + clean_mask = is_real_q & is_real_k & (kv_idx < current_block_end) + eye_mask = q_idx == kv_idx + return eye_mask | clean_mask + + block_mask = create_block_mask( + attention_mask, + B=batch_size, + H=None, + Q_LEN=total_length + padded_length, + KV_LEN=total_length + padded_length, + _compile=True, + device=device, + ) + + import torch.distributed as dist + if not dist.is_initialized() or dist.get_rank() == 0: + print( + f" cache a block wise causal mask with block size of {num_frame_per_block} frames" + ) + print(block_mask) + + return block_mask + + @staticmethod + def _prepare_teacher_forcing_mask( + device: torch.device | str, + num_frames: int = 31, + frame_seqlen: int = 880, + num_frame_per_block: int = 1, + batch_size: int | None = None, + ): + total_length = num_frames * frame_seqlen * 2 + padded_length = math.ceil(total_length / 128) * 128 - total_length + + clean_ends = num_frames * frame_seqlen + attention_block_size = frame_seqlen * num_frame_per_block + + # Use pure mathematical coordinates; do not introduce external tensor + # lookup tables here. + def attention_mask(b, h, q_idx, kv_idx): + is_real_q = q_idx < total_length + is_real_k = kv_idx < total_length + + # ========================================== + # 1. Clean-frame mask. + # ========================================== + is_clean_q = q_idx < clean_ends + + # End position of the block containing the current token. + clean_block_idx = q_idx // attention_block_size + current_clean_block_end = (clean_block_idx + 1) * attention_block_size + + clean_mask = ( + is_clean_q + & (kv_idx < current_clean_block_end) + ) + + # ========================================== + # 2. Noisy-frame mask. + # ========================================== + is_noisy_q = q_idx >= clean_ends + + noisy_rel_idx = q_idx - clean_ends + block_index = noisy_rel_idx // attention_block_size + + # C1: noisy tokens in the same block. + noisy_block_start = clean_ends + (block_index * attention_block_size) + noisy_block_end = noisy_block_start + attention_block_size + C1 = (kv_idx >= noisy_block_start) & (kv_idx < noisy_block_end) + + # C2: clean context tokens from previous blocks. + context_end_for_noisy = block_index * attention_block_size + + C2 = kv_idx < context_end_for_noisy + noise_mask = is_noisy_q & (C1 | C2) + + # ========================================== + # 3. Final merge. + # ========================================== + eye_mask = q_idx == kv_idx + return eye_mask | (is_real_q & is_real_k & (clean_mask | noise_mask)) + + # _compile=True is required here. Triton compiles the mathematical + # formula above directly into a memory-efficient block-sparse matrix. + block_mask = create_block_mask( + attention_mask, + B=batch_size, + H=None, + Q_LEN=total_length + padded_length, + KV_LEN=total_length + padded_length, + _compile=True, + device=device, + ) + + import torch.distributed as dist + if not dist.is_initialized() or dist.get_rank() == 0: + print(block_mask) + + return block_mask + + @staticmethod + def _prepare_teacher_forcing_mask_natural( + device: torch.device | str, + num_frames: int, + frame_seqlen: int, + num_frame_per_block: int = 1, + sp_size: int = 1, + batch_size: int | None = None, + ): + """Teacher-Forcing attention mask for the *natural* interleaved layout + produced directly by `all_to_all(scatter=head, gather=seq)`: + + [r0_clean, r0_noisy, r1_clean, r1_noisy, ..., r_{sp-1}_clean, r_{sp-1}_noisy] + + Semantically equivalent to :func:`_prepare_teacher_forcing_mask` (which + assumes the [all_clean; all_noisy] layout), but the mask decodes + ``is_noisy`` and ``global_frame`` directly from the token index, so + :func:`distributed_flex_attention` no longer has to reshape/permute + tokens after all_to_all. + """ + assert num_frames % sp_size == 0, ( + f"num_frames ({num_frames}) must be divisible by sp_size ({sp_size}) " + f"for natural TF layout" + ) + + F_local = num_frames // sp_size + clean_half = F_local * frame_seqlen # per-rank, clean side + per_rank_len = 2 * clean_half # per-rank, clean + noisy + total_length = num_frames * frame_seqlen * 2 + padded_length = math.ceil(total_length / 128) * 128 - total_length + + def attention_mask(b, h, q_idx, kv_idx): + is_real_q = q_idx < total_length + is_real_k = kv_idx < total_length + + # ---- decode q ---- + r_q = q_idx // per_rank_len + in_rank_q = q_idx % per_rank_len + is_noisy_q = in_rank_q >= clean_half + side_q = in_rank_q % clean_half # offset within clean/noisy half + global_f_q = r_q * F_local + side_q // frame_seqlen + block_q = global_f_q // num_frame_per_block + + # ---- decode k ---- + r_k = kv_idx // per_rank_len + in_rank_k = kv_idx % per_rank_len + is_noisy_k = in_rank_k >= clean_half + side_k = in_rank_k % clean_half + global_f_k = r_k * F_local + side_k // frame_seqlen + block_k = global_f_k // num_frame_per_block + + # 1. clean_q -> clean_k: blockwise causal. + clean2clean = ( + (~is_noisy_q) & (~is_noisy_k) + & (block_k <= block_q) + ) + + # 2. noisy_q -> clean_k: strictly earlier blocks. + noisy2clean = ( + is_noisy_q & (~is_noisy_k) + & (block_k < block_q) + ) + + # 3. noisy_q -> noisy_k: only tokens within the same block. + noisy2noisy = ( + is_noisy_q & is_noisy_k + & (block_k == block_q) + ) + + eye_mask = q_idx == kv_idx + return eye_mask | ( + is_real_q & is_real_k + & (clean2clean | noisy2clean | noisy2noisy) + ) + + block_mask = create_block_mask( + attention_mask, + B=batch_size, + H=None, + Q_LEN=total_length + padded_length, + KV_LEN=total_length + padded_length, + _compile=True, + device=device, + ) + + import torch.distributed as dist + if not dist.is_initialized() or dist.get_rank() == 0: + print( + f"[TF mask natural] sp_size={sp_size} F_local={F_local} " + f"clean_half={clean_half} per_rank_len={per_rank_len} " + f"total_length={total_length} " + f"num_frame_per_block={num_frame_per_block}" + ) + print(block_mask) + + return block_mask + + def _apply_cache_updates(self, kv_cache, cache_update_infos): + """ + Applies cache updates collected from multiple blocks. + Args: + kv_cache: List of cache dictionaries for each block + cache_update_infos: List of (block_index, cache_update_info) tuples + """ + for block_index, (current_end, local_end_index, update_info) in cache_update_infos: + if update_info is not None: + cache = kv_cache[block_index] + is_quantized = cache.get("quantized", False) + + if update_info["action"] == "roll_and_insert": + # Apply the rolling update. + sink_tokens = update_info["sink_tokens"] + num_rolled_tokens = update_info["num_rolled_tokens"] + num_evicted_tokens = update_info["num_evicted_tokens"] + local_start_index = update_info["local_start_index"] + local_end_index = update_info["local_end_index"] + new_k = update_info["new_k"] + new_v = update_info["new_v"] + + if is_quantized: + from utils.quant import copy_quantized_into, quantize_to_fp4 + + blk_sz = int(cache["block_token_size"]) + sink_blks = sink_tokens // blk_sz + evict_blks = num_evicted_tokens // blk_sz + roll_blks = num_rolled_tokens // blk_sz + + # iter-26: in-place copy into pre-allocated cache + # slots instead of replacing the QuantizedTensor + # reference. Required to unblock cudagraphs (the + # fresh QT returned by quantize_to_fp4 lives in + # the cudagraph memory pool; copying its data into + # the persistent slot buffer breaks that escape). + for i in range(roll_blks): + src = sink_blks + evict_blks + i + dst = sink_blks + i + copy_quantized_into(cache["k"][dst], cache["k"][src]) + copy_quantized_into(cache["v"][dst], cache["v"][src]) + + start_blk = local_start_index // blk_sz + n_insert_blks = (local_end_index - local_start_index) // blk_sz + head_dim = new_k.shape[-1] + for bi in range(n_insert_blks): + blk_idx = start_blk + bi + ts = bi * blk_sz + te = ts + blk_sz + k_block = new_k[0, ts:te, :, :].reshape(-1, head_dim).contiguous() + v_block = new_v[0, ts:te, :, :].reshape(-1, head_dim).contiguous() + copy_quantized_into( + cache["k"][blk_idx], + quantize_to_fp4(k_block, self.kv_quant_config), + ) + copy_quantized_into( + cache["v"][blk_idx], + quantize_to_fp4(v_block, self.kv_quant_config), + ) + else: + # Roll cached tokens. + cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ + cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() + cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ + cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() + + # Insert the new key/value tensors. + cache["k"][:, local_start_index:local_end_index] = new_k + cache["v"][:, local_start_index:local_end_index] = new_v + + # If a pinned multi-shot sink lives outside position 0, + # the rolling shifted everything left by num_evicted_tokens; + # pinned_start must follow so it tracks the same data. + pinned_shift = update_info.get("pinned_shift", 0) + if pinned_shift > 0 and "pinned_start" in cache: + cache["pinned_start"].sub_(pinned_shift) + + elif update_info["action"] == "direct_insert": + # Insert directly. + local_start_index = update_info["local_start_index"] + local_end_index = update_info["local_end_index"] + new_k = update_info["new_k"] + new_v = update_info["new_v"] + if is_quantized: + from utils.quant import copy_quantized_into, quantize_to_fp4 + + blk_sz = int(cache["block_token_size"]) + start_blk = local_start_index // blk_sz + n_insert_blks = (local_end_index - local_start_index) // blk_sz + head_dim = new_k.shape[-1] + # iter-26: in-place copy (see note above). + for bi in range(n_insert_blks): + blk_idx = start_blk + bi + ts = bi * blk_sz + te = ts + blk_sz + k_block = new_k[0, ts:te, :, :].reshape(-1, head_dim).contiguous() + v_block = new_v[0, ts:te, :, :].reshape(-1, head_dim).contiguous() + copy_quantized_into( + cache["k"][blk_idx], + quantize_to_fp4(k_block, self.kv_quant_config), + ) + copy_quantized_into( + cache["v"][blk_idx], + quantize_to_fp4(v_block, self.kv_quant_config), + ) + else: + # Insert the new key/value tensors. + cache["k"][:, local_start_index:local_end_index] = new_k + cache["v"][:, local_start_index:local_end_index] = new_v + + # Update cache indices. + kv_cache[block_index]["global_end_index"].fill_(current_end) + kv_cache[block_index]["local_end_index"].fill_(local_end_index) + + def _forward_inference( + self, + x, + t, + context, + seq_len, + clip_fea=None, + y=None, + kv_cache: dict = None, + crossattn_cache: dict = None, + current_start: int = 0, + cache_start: int = 0, + defer_cache_updates: bool = False, + ): + r""" + Run the diffusion model with kv caching. + See Algorithm 2 of CausVid paper https://arxiv.org/abs/2412.07772 for details. + This function will be run for num_frame times. + Process the latent frames one by one (880 tokens each) + + Args: + x (List[Tensor]): + List of input video tensors, each with shape [C_in, F, H, W] + t (Tensor): + Diffusion timesteps tensor of shape [B, F] + context (List[Tensor]): + List of text embeddings each with shape [L, C] + seq_len (`int`): + Maximum sequence length for positional encoding + clip_fea (Tensor, *optional*): + CLIP image features for image-to-video mode + y (List[Tensor], *optional*): + Conditional video inputs for image-to-video mode, same shape as x + + Returns: + List[Tensor]: + List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8] + """ + + if self.model_type == 'i2v': + assert clip_fea is not None and y is not None + # params + device = self.patch_embedding.weight.device + if self.freqs.device != device: + self.freqs = self.freqs.to(device) + + if y is not None: + x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)] + + # embeddings + x = [self.patch_embedding(u.unsqueeze(0)) for u in x] + # iter-31: publish chunk metadata as Python ints to module-level dict + # so deep attention forwards can avoid `.item()` graph breaks. + first_shape = tuple(x[0].shape[2:]) + _CURRENT_GRID_META["frame_seqlen"] = int(first_shape[1] * first_shape[2]) + _CURRENT_GRID_META["num_new_frames"] = int(first_shape[0]) + _CURRENT_GRID_META["h"] = int(first_shape[1]) + _CURRENT_GRID_META["w"] = int(first_shape[2]) + # iter-39 v2: kv_cache scalars (global_end_index, local_end_index, + # pinned_start, pinned_len) are published into _CURRENT_GRID_META by + # the eager `_call_model` wrapper (utils/wan_5b_wrapper.py) BEFORE + # this compiled forward runs. Reading them here via `.item()` would + # trigger graph breaks; the wrapper does it in eager Python instead. + grid_sizes = torch.stack( + [torch.tensor(u.shape[2:], dtype=torch.long) for u in x]) + x = [u.flatten(2).transpose(1, 2) for u in x] + seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long) + x = torch.cat(x) + + # time embeddings + if t.dim() == 1: + raise NotImplementedError(f"t.shape should be [B, F], but got {t.shape}") + + bt = t.size(0) + t_len = t.size(1) + t = t.flatten() + e = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, + t).unflatten(0, (bt, t_len)).type_as(x)) + e0 = self.time_projection(e).unflatten(2, (6, self.dim)) # B, F, 6, C + + # context + context_lens = None + context = self.text_embedding( + torch.stack([ + torch.cat( + [u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) + for u in context + ])) + + # arguments + kwargs = dict( + e=e0, + seq_lens=seq_lens, + grid_sizes=grid_sizes, + freqs=self.freqs, + context=context, + context_lens=context_lens, + block_mask=self.block_mask, + t_scale=self.t_scale, + use_relative_rope=self.use_relative_rope, + method=self.rope_method, + original_seq_len=self.original_seq_len, + temporal_offset=self.rope_temporal_offset, + ) + + def create_custom_forward(module): + def custom_forward(*inputs, **kwargs): + return module(*inputs, **kwargs) + return custom_forward + + cache_update_info = None + cache_update_infos = [] # Collect cache updates from every block. + for block_index, block in enumerate(self.blocks): + if torch.is_grad_enabled() and self.gradient_checkpointing: + kwargs.update( + { + "kv_cache": kv_cache[block_index], + "current_start": current_start, + "cache_start": cache_start + } + ) + result = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, **kwargs, + use_reentrant=False, + ) + # Handle the result + if kv_cache is not None and isinstance(result, tuple): + x, block_cache_update_info = result + cache_update_infos.append((block_index, block_cache_update_info)) + # Keep only basic metadata for later blocks, without the + # concrete cache-update payload. + cache_update_info = block_cache_update_info[:2] # (current_end, local_end_index) + else: + x = result + else: + kwargs.update( + { + "kv_cache": kv_cache[block_index], + "crossattn_cache": crossattn_cache[block_index], + "current_start": current_start, + "cache_start": cache_start + } + ) + result = block(x, **kwargs) + # Handle the result + if kv_cache is not None and isinstance(result, tuple): + x, block_cache_update_info = result + cache_update_infos.append((block_index, block_cache_update_info)) + # Keep only basic metadata for later blocks, without the + # concrete cache-update payload. + cache_update_info = block_cache_update_info[:2] # (current_end, local_end_index) + else: + x = result + + # Apply all cache updates after every block has run. For cudagraphs + # experiments this can be deferred to the eager wrapper so cache + # mutation does not happen inside the compiled forward. + if kv_cache is not None and cache_update_infos and not defer_cache_updates: + self._apply_cache_updates(kv_cache, cache_update_infos) + + # head + x = self.head(x, e.unsqueeze(2)) + # unpatchify + x = self.unpatchify(x, grid_sizes) + output = torch.stack(x) + if kv_cache is not None and defer_cache_updates: + return output, cache_update_infos + return output + + def _forward_train( + self, + x, + t, + context, + seq_len, + clean_x=None, + aug_t=None, + clip_fea=None, + y=None, + ): + r""" + Forward pass through the diffusion model + + Args: + x (List[Tensor]): + List of input video tensors, each with shape [C_in, F, H, W] + t (Tensor): + Diffusion timesteps tensor of shape [B] + context (List[Tensor]): + List of text embeddings each with shape [L, C] + seq_len (`int`): + Maximum sequence length for positional encoding + clip_fea (Tensor, *optional*): + CLIP image features for image-to-video mode + y (List[Tensor], *optional*): + Conditional video inputs for image-to-video mode, same shape as x + + Returns: + List[Tensor]: + List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8] + """ + if self.model_type == 'i2v': + assert clip_fea is not None and y is not None + # params + device = self.patch_embedding.weight.device + if self.freqs.device != device: + self.freqs = self.freqs.to(device) + + # Construct blockwise causal attn mask + # Recreate mask when batch size changes to avoid Triton broadcasting bug + current_batch_size = x.shape[0] + if self.block_mask is None or self._block_mask_batch_size != current_batch_size: + self._block_mask_batch_size = current_batch_size + if clean_x is not None: + if self.independent_first_frame: + raise NotImplementedError() + else: + self.block_mask = self._prepare_teacher_forcing_mask( + device, num_frames=x.shape[2], + frame_seqlen=x.shape[-2] * x.shape[-1] // (self.patch_size[1] * self.patch_size[2]), + num_frame_per_block=self.num_frame_per_block, + batch_size=current_batch_size, + ) + else: + if self.independent_first_frame: + self.block_mask = self._prepare_blockwise_causal_attn_mask_i2v( + device, num_frames=x.shape[2], + frame_seqlen=x.shape[-2] * x.shape[-1] // (self.patch_size[1] * self.patch_size[2]), + num_frame_per_block=self.num_frame_per_block, + batch_size=current_batch_size, + ) + else: + self.block_mask = self._prepare_blockwise_causal_attn_mask( + device, num_frames=x.shape[2], + frame_seqlen=x.shape[-2] * x.shape[-1] // (self.patch_size[1] * self.patch_size[2]), + num_frame_per_block=self.num_frame_per_block, + batch_size=current_batch_size, + ) + + if y is not None: + x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)] + + # embeddings + x = [self.patch_embedding(u.unsqueeze(0)) for u in x] + + grid_sizes = torch.stack( + [torch.tensor(u.shape[2:], dtype=torch.long) for u in x]) + x = [u.flatten(2).transpose(1, 2) for u in x] + + seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long) + max_len = int(seq_lens.max().item()) + assert max_len > 0, "Token sequence length is zero after patch embedding" + # Pad all samples to the batch max length instead of the first sample length + x = torch.cat([ + torch.cat([u, u.new_zeros(1, max_len - u.size(1), u.size(2))], dim=1) + for u in x + ]) + + # time embeddings + if t.dim() == 1: + raise NotImplementedError(f"t.shape should be [B, F], but got {t.shape}") + bt = t.size(0) + t_len = t.size(1) + t_ori_shape = t.shape + t = t.flatten() + e = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, t).unflatten(0, (bt, t_len)).type_as(x)) + e0 = self.time_projection(e).unflatten(2, (6, self.dim)) # B, F, 6, C + + # context + context_lens = None + context = self.text_embedding( + torch.stack([ + torch.cat( + [u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) + for u in context + ])) + + if clean_x is not None: + clean_x = [self.patch_embedding(u.unsqueeze(0)) for u in clean_x] + clean_x = [u.flatten(2).transpose(1, 2) for u in clean_x] + + seq_lens_clean = torch.tensor([u.size(1) for u in clean_x], dtype=torch.long) + clean_x = torch.cat([ + torch.cat([u, u.new_zeros(1, seq_lens_clean[0] - u.size(1), u.size(2))], dim=1) for u in clean_x + ]) + + x = torch.cat([clean_x, x], dim=1) + if aug_t is None: + aug_t = torch.zeros(t_ori_shape, device=t.device, dtype=t.dtype) + bt_clean = aug_t.size(0) + t_clean_len = aug_t.size(1) + aug_t = aug_t.flatten() + e_clean = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, aug_t).unflatten(0, (bt_clean, t_clean_len)).type_as(x)) + e0_clean = self.time_projection(e_clean).unflatten(2, (6, self.dim)) + e0 = torch.cat([e0_clean, e0], dim=1) + + # arguments + kwargs = dict( + e=e0, + seq_lens=seq_lens, + grid_sizes=grid_sizes, + freqs=self.freqs, + context=context, + context_lens=context_lens, + block_mask=self.block_mask, + t_scale=self.t_scale, + method=self.rope_method, + original_seq_len=self.original_seq_len, + temporal_offset=self.rope_temporal_offset, + ) + + def create_custom_forward(module): + def custom_forward(*inputs, **kwargs): + return module(*inputs, **kwargs) + return custom_forward + + for block in self.blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, **kwargs, + use_reentrant=False, + ) + else: + x = block(x, **kwargs) + + if clean_x is not None: + x = x[:, x.shape[1] // 2:] + + # head + x = self.head(x, e.unsqueeze(2)) + + # unpatchify + x = self.unpatchify(x, grid_sizes) + return torch.stack(x) + + def forward( + self, + *args, + **kwargs + ): + if kwargs.get('kv_cache', None) is not None: + return self._forward_inference(*args, **kwargs) + else: + return self._forward_train(*args, **kwargs) + + def unpatchify(self, x, grid_sizes): + r""" + Reconstruct video tensors from patch embeddings. + + Args: + x (List[Tensor]): + List of patchified features, each with shape [L, C_out * prod(patch_size)] + grid_sizes (Tensor): + Original spatial-temporal grid dimensions before patching, + shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches) + + Returns: + List[Tensor]: + Reconstructed video tensors with shape [C_out, F, H / 8, W / 8] + """ + + c = self.out_dim + out = [] + for u, v in zip(x, grid_sizes.tolist()): + u = u[:math.prod(v)].view(*v, *self.patch_size, c) + u = torch.einsum('fhwpqrc->cfphqwr', u) + u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)]) + out.append(u) + return out + + def init_weights(self): + r""" + Initialize model parameters using Xavier initialization. + """ + + # basic init + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + if m.bias is not None: + nn.init.zeros_(m.bias) + + # init embeddings + nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1)) + for m in self.text_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=.02) + for m in self.time_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=.02) + + # init output layer + nn.init.zeros_(self.head.head.weight) diff --git a/wan_5b/modules/causal_model_sp_ulysses.py b/wan_5b/modules/causal_model_sp_ulysses.py new file mode 100644 index 0000000..7d98bb9 --- /dev/null +++ b/wan_5b/modules/causal_model_sp_ulysses.py @@ -0,0 +1,601 @@ +# Copyright 2024-2025 LongLive Authors. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +""" +Causal Wan2.2-TI2V-5B model with Ulysses-style sequence parallelism. + +This inference-only variant mirrors the regular CausalWanModel parameter layout +so existing generator checkpoints can be loaded into ``model.*`` keys. +""" + +import math +import os + +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin + +from wan_5b.modules.attention import attention +from wan_5b.modules.model import ( + WanRMSNorm, + WanLayerNorm, + rope_apply, + rope_params, + sinusoidal_embedding_1d, +) +from wan_5b.modules.causal_model import causal_rope_apply, MultiShotT2VCrossAttention +from wan_5b.distributed.sp_ulysses_inference import ( + get_sp_rank, + get_sp_world_size, + is_sp_enabled, + sp_all_gather, + sp_scatter, + ulysses_head_to_seq, + ulysses_seq_to_head, +) + + +QK_QUANT_SIM = os.environ.get("QK_QUANT_SIM", "0") == "1" +QK_QUANT_RATIO = float(os.environ.get("QK_QUANT_RATIO", "0.28125")) + + +try: + import torch.cuda.nvtx as nvtx + NVTX_ENABLED = True +except Exception: + nvtx = None + NVTX_ENABLED = False + + +class NVTXRange: + def __init__(self, name: str): + self.name = name + + def __enter__(self): + if NVTX_ENABLED: + nvtx.range_push(self.name) + + def __exit__(self, exc_type, exc, tb): + if NVTX_ENABLED: + nvtx.range_pop() + + +def _get_d_quant(head_dim: int, ratio: float = QK_QUANT_RATIO) -> int: + d_quant = max(8, round(head_dim * ratio)) + return ((d_quant + 7) // 8) * 8 + + +def _compute_ulysses_frame_info(num_frames, local_seq_len, sp_world_size, sp_rank): + if num_frames == 1: + return 1, local_seq_len, 0 + assert num_frames % sp_world_size == 0, ( + f"Ulysses SP requires num_frames ({num_frames}) divisible by " + f"world_size ({sp_world_size})." + ) + local_frames = num_frames // sp_world_size + frame_seqlen = local_seq_len // local_frames + frame_offset = sp_rank * local_frames + return local_frames, frame_seqlen, frame_offset + + +class UlyssesCausalWanSelfAttention(nn.Module): + """Causal self-attention with Ulysses sequence/head exchange.""" + + def __init__(self, dim, num_heads, local_attn_size=-1, sink_size=0, + qk_norm=True, eps=1e-6): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.local_attn_size = local_attn_size + self.sink_size = sink_size + self.global_sink_size = 0 + self.qk_norm = qk_norm + self.eps = eps + + if not isinstance(local_attn_size, int) and hasattr(local_attn_size, "__iter__"): + values = list(local_attn_size) + else: + values = [int(local_attn_size)] + non_neg_vals = [int(v) for v in values if int(v) != -1] + max_local = max(non_neg_vals) if non_neg_vals else -1 + self.max_attention_size = 28160 if max_local == -1 else max_local * 880 + + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + + def forward(self, x, seq_lens, grid_sizes, freqs, kv_cache=None, + current_start=0, cache_start=None, t_scale=1.0, + use_relative_rope=False, method="linear", + original_seq_len=None, temporal_offset=0.0): + b, s_local = x.shape[:2] + n, d = self.num_heads, self.head_dim + if cache_start is None: + cache_start = current_start + + with NVTXRange("ulysses_qkv_proj"): + q = self.norm_q(self.q(x)).view(b, s_local, n, d) + k = self.norm_k(self.k(x)).view(b, s_local, n, d) + v = self.v(x).view(b, s_local, n, d) + + if kv_cache is None: + return self._forward_no_cache(q, k, v, grid_sizes, freqs, t_scale, method, original_seq_len) + return self._forward_with_cache( + q, k, v, grid_sizes, freqs, kv_cache, current_start, cache_start, + t_scale, use_relative_rope, method, original_seq_len, temporal_offset, + ) + + def _exchange_qkv(self, q, k, v): + if QK_QUANT_SIM and is_sp_enabled(): + d = q.shape[-1] + d_quant = _get_d_quant(d) + q_heads = F.pad(ulysses_seq_to_head(q[..., :d_quant]), (0, d - d_quant)) + k_heads = F.pad(ulysses_seq_to_head(k[..., :d_quant]), (0, d - d_quant)) + v_heads = ulysses_seq_to_head(v) + else: + q_heads = ulysses_seq_to_head(q) + k_heads = ulysses_seq_to_head(k) + v_heads = ulysses_seq_to_head(v) + return q_heads, k_heads, v_heads + + def _forward_no_cache(self, q, k, v, grid_sizes, freqs, t_scale, method, original_seq_len): + with NVTXRange("ulysses_seq_to_head"): + q_heads, k_heads, v_heads = self._exchange_qkv(q, k, v) + + with NVTXRange("ulysses_rope"): + roped_q = rope_apply( + q_heads, grid_sizes, freqs, t_scale=t_scale, + method=method, original_seq_len=original_seq_len, + ).type_as(v_heads) + roped_k = rope_apply( + k_heads, grid_sizes, freqs, t_scale=t_scale, + method=method, original_seq_len=original_seq_len, + ).type_as(v_heads) + + with NVTXRange("ulysses_attention"): + out_heads = attention(roped_q, roped_k, v_heads, causal=True) + with NVTXRange("ulysses_head_to_seq"): + out = ulysses_head_to_seq(out_heads) + with NVTXRange("ulysses_out_proj"): + return self.o(out.flatten(2)) + + def _forward_with_cache(self, q, k, v, grid_sizes, freqs, kv_cache, + current_start, cache_start, t_scale, + use_relative_rope, method, original_seq_len, + temporal_offset): + sp_world_size = get_sp_world_size() if is_sp_enabled() else 1 + s_total = q.shape[1] * sp_world_size + frame_seqlen = math.prod(grid_sizes[0][1:]).item() + current_start_frame = current_start // frame_seqlen + current_end = current_start + s_total + + with NVTXRange("ulysses_seq_to_head"): + q_heads, k_heads, v_heads = self._exchange_qkv(q, k, v) + + with NVTXRange("ulysses_rope_cached"): + roped_q = causal_rope_apply( + q_heads, grid_sizes, freqs, start_frame=current_start_frame, + t_scale=t_scale, method=method, original_seq_len=original_seq_len, + temporal_offset=temporal_offset, + ).type_as(v_heads) + if use_relative_rope: + key_to_cache = k_heads + else: + key_to_cache = causal_rope_apply( + k_heads, grid_sizes, freqs, start_frame=current_start_frame, + t_scale=t_scale, method=method, original_seq_len=original_seq_len, + temporal_offset=temporal_offset, + ).type_as(v_heads) + + with NVTXRange("ulysses_cache_update"): + k_full, v_full = self._update_cache_and_get_kv( + key_to_cache, v_heads, kv_cache, current_start, current_end, frame_seqlen, + ) + + if use_relative_rope: + raise NotImplementedError("use_relative_rope is not implemented for SP inference.") + + with NVTXRange("ulysses_cached_attention"): + out_heads = attention(roped_q, k_full, v_full, causal=False) + with NVTXRange("ulysses_head_to_seq"): + out = ulysses_head_to_seq(out_heads) + with NVTXRange("ulysses_out_proj"): + out = self.o(out.flatten(2)) + return out, (current_end, kv_cache["local_end_index"].item(), None) + + def _effective_sink(self, kv_cache, frame_seqlen): + sink_tokens = self.sink_size * frame_seqlen + global_sink_tokens = getattr(self, "global_sink_size", 0) * frame_seqlen + pinned_start_t = kv_cache.get("pinned_start", None) + if pinned_start_t is not None and hasattr(pinned_start_t, "item"): + pinned_start = pinned_start_t.item() + pinned_len = kv_cache["pinned_len"].item() + else: + pinned_start = -1 + pinned_len = 0 + has_pinned = pinned_start >= 0 and pinned_len > 0 + if has_pinned and pinned_start == global_sink_tokens: + effective_sink = global_sink_tokens + pinned_len + elif has_pinned: + effective_sink = global_sink_tokens + else: + effective_sink = max(global_sink_tokens, sink_tokens) + return effective_sink, pinned_start, pinned_len, has_pinned + + def _update_cache_and_get_kv(self, k_new, v_new, kv_cache, current_start, + current_end, frame_seqlen): + b, s_new, _, _ = k_new.shape + kv_cache_size = kv_cache["k"].shape[1] + global_end_prev = kv_cache["global_end_index"].item() + local_end_prev = kv_cache["local_end_index"].item() + is_recompute = current_end <= global_end_prev and current_start > 0 + + effective_sink, pinned_start, pinned_len, has_pinned = self._effective_sink(kv_cache, frame_seqlen) + need_roll = ( + self.local_attn_size != -1 + and current_end > global_end_prev + and s_new + local_end_prev > kv_cache_size + ) + + if need_roll: + num_evicted = s_new + local_end_prev - kv_cache_size + num_rolled = local_end_prev - num_evicted - effective_sink + if num_rolled > 0: + kv_cache["k"][:, effective_sink:effective_sink + num_rolled] = ( + kv_cache["k"][ + :, effective_sink + num_evicted:effective_sink + num_evicted + num_rolled + ].clone() + ) + kv_cache["v"][:, effective_sink:effective_sink + num_rolled] = ( + kv_cache["v"][ + :, effective_sink + num_evicted:effective_sink + num_evicted + num_rolled + ].clone() + ) + if has_pinned and pinned_start >= effective_sink and "pinned_start" in kv_cache: + kv_cache["pinned_start"].sub_(num_evicted) + pinned_start -= num_evicted + local_end_new = local_end_prev + (current_end - global_end_prev) - num_evicted + else: + local_end_new = local_end_prev + (current_end - global_end_prev) + + local_start_new = local_end_new - s_new + write_start = max(local_start_new, effective_sink) if is_recompute else local_start_new + write_offset = max(0, write_start - local_start_new) + write_len = max(0, local_end_new - write_start) + if write_len > 0: + kv_cache["k"][:, write_start:local_end_new] = k_new[:, write_offset:write_offset + write_len] + kv_cache["v"][:, write_start:local_end_new] = v_new[:, write_offset:write_offset + write_len] + + if not is_recompute: + kv_cache["global_end_index"].fill_(current_end) + kv_cache["local_end_index"].fill_(local_end_new) + + window_start = max(0, local_end_new - self.max_attention_size) + prepend_sink = effective_sink > 0 and window_start > 0 + prepend_pinned = has_pinned and pinned_start >= effective_sink and pinned_start < window_start + + if prepend_sink and prepend_pinned: + extra = effective_sink + pinned_len + effective_local_size = max(0, self.max_attention_size - extra) + local_window_start = max(effective_sink, local_end_new - effective_local_size) + k_full = torch.cat([ + kv_cache["k"][:, :effective_sink], + kv_cache["k"][:, pinned_start:pinned_start + pinned_len], + kv_cache["k"][:, local_window_start:local_end_new], + ], dim=1) + v_full = torch.cat([ + kv_cache["v"][:, :effective_sink], + kv_cache["v"][:, pinned_start:pinned_start + pinned_len], + kv_cache["v"][:, local_window_start:local_end_new], + ], dim=1) + elif prepend_sink: + effective_local_size = max(0, self.max_attention_size - effective_sink) + local_window_start = max(effective_sink, local_end_new - effective_local_size) + k_full = torch.cat([ + kv_cache["k"][:, :effective_sink], + kv_cache["k"][:, local_window_start:local_end_new], + ], dim=1) + v_full = torch.cat([ + kv_cache["v"][:, :effective_sink], + kv_cache["v"][:, local_window_start:local_end_new], + ], dim=1) + elif prepend_pinned: + effective_local_size = max(0, self.max_attention_size - pinned_len) + local_window_start = max(0, local_end_new - effective_local_size) + k_full = torch.cat([ + kv_cache["k"][:, pinned_start:pinned_start + pinned_len], + kv_cache["k"][:, local_window_start:local_end_new], + ], dim=1) + v_full = torch.cat([ + kv_cache["v"][:, pinned_start:pinned_start + pinned_len], + kv_cache["v"][:, local_window_start:local_end_new], + ], dim=1) + else: + k_full = kv_cache["k"][:, window_start:local_end_new] + v_full = kv_cache["v"][:, window_start:local_end_new] + + return k_full, v_full + + +class UlyssesCausalWanAttentionBlock(nn.Module): + def __init__(self, dim, ffn_dim, num_heads, local_attn_size=-1, + sink_size=0, qk_norm=True, cross_attn_norm=False, eps=1e-6): + super().__init__() + self.norm1 = WanLayerNorm(dim, eps) + self.self_attn = UlyssesCausalWanSelfAttention( + dim, num_heads, local_attn_size, sink_size, qk_norm, eps + ) + self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() + self.cross_attn = MultiShotT2VCrossAttention(dim, num_heads, (-1, -1), qk_norm, eps) + self.norm2 = WanLayerNorm(dim, eps) + self.ffn = nn.Sequential( + nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim) + ) + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward(self, x, e, seq_lens, grid_sizes, freqs, context, context_lens, + kv_cache=None, crossattn_cache=None, current_start=0, + cache_start=None, num_frames=None, local_frame_offset=0, + t_scale=1.0, use_relative_rope=False, method="linear", + original_seq_len=None, temporal_offset=0.0): + num_frames_total = e.shape[1] if num_frames is None else num_frames + local_seq_len = x.shape[1] + if is_sp_enabled() and local_seq_len > 0: + actual_local_frames, local_frame_seqlen, _ = _compute_ulysses_frame_info( + num_frames_total, local_seq_len, get_sp_world_size(), get_sp_rank() + ) + e_local = e[:, local_frame_offset:local_frame_offset + actual_local_frames] + else: + actual_local_frames = num_frames_total + local_frame_seqlen = local_seq_len // actual_local_frames if actual_local_frames > 0 else 0 + e_local = e + + e_chunks = (self.modulation.unsqueeze(1) + e_local).chunk(6, dim=2) + x_normed = self.norm1(x) + if actual_local_frames > 0 and local_frame_seqlen > 0: + x_mod = x_normed.unflatten(dim=1, sizes=(actual_local_frames, local_frame_seqlen)) + x_mod = (x_mod * (1 + e_chunks[1]) + e_chunks[0]).flatten(1, 2) + else: + x_mod = x_normed + + self_attn_result = self.self_attn( + x_mod, seq_lens, grid_sizes, freqs, kv_cache, current_start, cache_start, + t_scale, use_relative_rope, method, original_seq_len, temporal_offset, + ) + if kv_cache is not None and isinstance(self_attn_result, tuple): + y, cache_update_info = self_attn_result + else: + y = self_attn_result + cache_update_info = None + + if actual_local_frames > 0 and local_frame_seqlen > 0: + y_mod = y.unflatten(dim=1, sizes=(actual_local_frames, local_frame_seqlen)) + x = x + (y_mod * e_chunks[2]).flatten(1, 2) + else: + x = x + y + + if x.shape[1] > 0: + x = x + self.cross_attn( + self.norm3(x), context, context_lens, crossattn_cache=crossattn_cache + ) + x_normed = self.norm2(x) + if actual_local_frames > 0 and local_frame_seqlen > 0: + x_mod = x_normed.unflatten(dim=1, sizes=(actual_local_frames, local_frame_seqlen)) + y = self.ffn((x_mod * (1 + e_chunks[4]) + e_chunks[3]).flatten(1, 2)) + y_mod = y.unflatten(dim=1, sizes=(actual_local_frames, local_frame_seqlen)) + x = x + (y_mod * e_chunks[5]).flatten(1, 2) + else: + x = x + self.ffn(x_normed) + + return (x, cache_update_info) if cache_update_info is not None else x + + +class UlyssesCausalHead(nn.Module): + def __init__(self, dim, out_dim, patch_size, eps=1e-6): + super().__init__() + self.dim = dim + self.out_dim = out_dim + self.patch_size = patch_size + self.norm = WanLayerNorm(dim, eps) + self.head = nn.Linear(dim, math.prod(patch_size) * out_dim) + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x, e, local_frame_offset=0, actual_local_frames=None, + frame_seqlen_global=None): + num_frames_total = e.shape[1] + local_seq_len = x.shape[1] + if is_sp_enabled() and actual_local_frames is not None: + local_frames = actual_local_frames + local_seqlen = local_seq_len // local_frames if local_frames > 0 else frame_seqlen_global + e_local = e[:, local_frame_offset:local_frame_offset + local_frames] + else: + local_frames = num_frames_total + local_seqlen = local_seq_len // local_frames if local_frames > 0 else 0 + e_local = e + + if local_frames == 0 or local_seq_len == 0: + b, _, _ = x.shape + return torch.empty( + b, 0, local_seqlen if local_seqlen else 1, + math.prod(self.patch_size) * self.out_dim, + dtype=x.dtype, device=x.device, + ) + + e_chunks = (self.modulation.unsqueeze(1) + e_local).chunk(2, dim=2) + x = self.norm(x).unflatten(dim=1, sizes=(local_frames, local_seqlen)) + return self.head(x * (1 + e_chunks[1]) + e_chunks[0]) + + +class UlyssesSPCausalWanModel(ModelMixin, ConfigMixin): + ignore_for_config = ["patch_size", "cross_attn_norm", "qk_norm", "text_dim"] + _no_split_modules = ["UlyssesCausalWanAttentionBlock"] + _supports_gradient_checkpointing = True + + @register_to_config + def __init__(self, model_type="ti2v", patch_size=(1, 2, 2), text_len=512, + in_dim=48, dim=3072, ffn_dim=14336, freq_dim=256, + text_dim=4096, out_dim=48, num_heads=24, num_layers=30, + local_attn_size=-1, sink_size=0, num_frame_per_block=1, + qk_norm=True, cross_attn_norm=True, eps=1e-6): + super().__init__() + assert model_type in ["t2v", "i2v", "ti2v"] + self.model_type = model_type + self.patch_size = patch_size + self.text_len = text_len + self.in_dim = in_dim + self.dim = dim + self.ffn_dim = ffn_dim + self.freq_dim = freq_dim + self.text_dim = text_dim + self.out_dim = out_dim + self.num_heads = num_heads + self.num_layers = num_layers + self.local_attn_size = local_attn_size + self.sink_size = sink_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + + self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, dim) + ) + self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6)) + self.blocks = nn.ModuleList([ + UlyssesCausalWanAttentionBlock( + dim, ffn_dim, num_heads, local_attn_size, sink_size, + qk_norm, cross_attn_norm, eps, + ) + for _ in range(num_layers) + ]) + self.head = UlyssesCausalHead(dim, out_dim, patch_size, eps) + + d = dim // num_heads + self.freqs = torch.cat([ + rope_params(1024, d - 4 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + ], dim=1) + + self.init_weights() + self.gradient_checkpointing = False + self.num_frame_per_block = num_frame_per_block + self.t_scale = 1.0 + self.use_relative_rope = False + self.rope_method = "linear" + self.original_seq_len = None + self.rope_temporal_offset = 0.0 + self.kv_quant_config = None + + def forward(self, x, t, context, seq_len, clip_fea=None, y=None, + kv_cache=None, crossattn_cache=None, current_start=0, + cache_start=0): + sp_world_size = get_sp_world_size() if is_sp_enabled() else 1 + sp_rank = get_sp_rank() if is_sp_enabled() else 0 + device = self.patch_embedding.weight.device + if self.freqs.device != device: + self.freqs = self.freqs.to(device) + if y is not None: + x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)] + + with NVTXRange("ulysses_patch_embedding"): + x = [self.patch_embedding(u.unsqueeze(0)) for u in x] + grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x]) + x = [u.flatten(2).transpose(1, 2) for u in x] + seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long) + assert seq_lens.max() <= seq_len + x = torch.cat(x) + + num_frames = t.shape[1] if t.dim() > 1 else grid_sizes[0, 0].item() + frame_seqlen = x.shape[1] // num_frames + if is_sp_enabled(): + assert self.num_heads % sp_world_size == 0 + assert x.shape[1] % sp_world_size == 0 + with NVTXRange("ulysses_scatter_input"): + x = sp_scatter(x, dim=1) + actual_local_frames, local_frame_seqlen, local_frame_offset = _compute_ulysses_frame_info( + num_frames, x.shape[1], sp_world_size, sp_rank + ) + else: + actual_local_frames = num_frames + local_frame_seqlen = frame_seqlen + local_frame_offset = 0 + + with NVTXRange("ulysses_time_embedding"): + e = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t.flatten()).type_as(x)) + e0 = self.time_projection(e).unflatten(1, (6, self.dim)).unflatten(dim=0, sizes=t.shape) + + with NVTXRange("ulysses_text_embedding"): + context_lens = None + context = self.text_embedding(torch.stack([ + torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) + for u in context + ])) + if clip_fea is not None: + context = torch.concat([self.img_emb(clip_fea), context], dim=1) + + kwargs = dict( + e=e0, + seq_lens=seq_lens, + grid_sizes=grid_sizes, + freqs=self.freqs, + context=context, + context_lens=context_lens, + num_frames=num_frames, + local_frame_offset=local_frame_offset, + t_scale=self.t_scale, + use_relative_rope=self.use_relative_rope, + method=self.rope_method, + original_seq_len=self.original_seq_len, + temporal_offset=self.rope_temporal_offset, + ) + for block_idx, block in enumerate(self.blocks): + kwargs.update({ + "kv_cache": kv_cache[block_idx] if kv_cache else None, + "crossattn_cache": crossattn_cache[block_idx] if crossattn_cache else None, + "current_start": current_start, + "cache_start": cache_start, + }) + result = block(x, **kwargs) + x = result[0] if kv_cache is not None and isinstance(result, tuple) else result + + x = self.head( + x, + e.unflatten(dim=0, sizes=t.shape).unsqueeze(2), + local_frame_offset=local_frame_offset, + actual_local_frames=actual_local_frames, + frame_seqlen_global=frame_seqlen, + ) + if is_sp_enabled(): + x = sp_all_gather(x, dim=2 if num_frames == 1 else 1) + x = self.unpatchify(x, grid_sizes) + return torch.stack(x) + + def unpatchify(self, x, grid_sizes): + c = self.out_dim + pt, ph, pw = self.patch_size + output = [] + for i, (f, h, w) in enumerate(grid_sizes.tolist()): + x_i = x[i, :f, :h * w, :].reshape(f, h, w, pt, ph, pw, c) + x_i = x_i.permute(6, 0, 3, 1, 4, 2, 5).reshape(c, f * pt, h * ph, w * pw) + output.append(x_i) + return output + + def init_weights(self): + def _init(module): + if isinstance(module, nn.Linear): + nn.init.xavier_uniform_(module.weight) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Conv3d): + nn.init.xavier_uniform_(module.weight) + if module.bias is not None: + nn.init.zeros_(module.bias) + self.apply(_init) diff --git a/wan_5b/modules/model.py b/wan_5b/modules/model.py new file mode 100644 index 0000000..b37b435 --- /dev/null +++ b/wan_5b/modules/model.py @@ -0,0 +1,591 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import math + +import torch +import torch.nn as nn +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +from einops import repeat +import torch.distributed as dist +from .attention import flash_attention + +__all__ = ['WanModel'] + + +def sinusoidal_embedding_1d(dim, position): + # preprocess + assert dim % 2 == 0 + half = dim // 2 + position = position.type(torch.float64) + + # calculation + sinusoid = torch.outer( + position, torch.pow(10000, -torch.arange(half).to(position).div(half))) + x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1) + return x + + +# @amp.autocast(enabled=False) +def rope_params(max_seq_len, dim, theta=10000): + assert dim % 2 == 0 + freqs = torch.outer( + torch.arange(max_seq_len), + 1.0 / torch.pow(theta, + torch.arange(0, dim, 2).to(torch.float64).div(dim))) + freqs = torch.polar(torch.ones_like(freqs), freqs) + return freqs + + +from utils.position_embedding_utils import ( + compute_temporal_freqs as _compute_temporal_freqs, + select_temporal_offset_for_sample, +) + + +# @amp.autocast(enabled=False) +def rope_apply(x, grid_sizes, freqs, t_scale=1.0, method="linear", original_seq_len=None, + temporal_offset=0.0): + n, c = x.size(2), x.size(3) // 2 + + # split freqs + freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) + + # loop over samples + output = [] + for i, (f, h, w) in enumerate(grid_sizes.tolist()): + seq_len = f * h * w + + # precompute multipliers + x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape( + seq_len, n, -1, 2)) + temporal_offset_i = select_temporal_offset_for_sample( + temporal_offset, i, f, start_frame=0) + temporal_freqs = _compute_temporal_freqs(freqs[0], f, 0, t_scale, x.device, + method=method, original_seq_len=original_seq_len, + temporal_offset=temporal_offset_i) + freqs_i = torch.cat([ + temporal_freqs.view(f, 1, 1, -1).expand(f, h, w, -1), + freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1) + ], + dim=-1).reshape(seq_len, 1, -1) + + # apply rotary embedding + x_i = torch.view_as_real(x_i * freqs_i).flatten(2) + x_i = torch.cat([x_i, x[i, seq_len:]]) + + # append to collection + output.append(x_i) + return torch.stack(output).type_as(x) + + +class WanRMSNorm(nn.Module): + + def __init__(self, dim, eps=1e-5): + super().__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + r""" + Args: + x(Tensor): Shape [B, L, C] + """ + return self._norm(x.float()).type_as(x) * self.weight + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) + + +class WanLayerNorm(nn.LayerNorm): + + def __init__(self, dim, eps=1e-6, elementwise_affine=False): + super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps) + + def forward(self, x): + r""" + Args: + x(Tensor): Shape [B, L, C] + """ + return super().forward(x).type_as(x) + + +class WanSelfAttention(nn.Module): + + def __init__(self, + dim, + num_heads, + window_size=(-1, -1), + qk_norm=True, + eps=1e-6): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.window_size = window_size + self.qk_norm = qk_norm + self.eps = eps + + # layers + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + + def forward(self, x, seq_lens, grid_sizes, freqs): + r""" + Args: + x(Tensor): Shape [B, L, num_heads, C / num_heads] + seq_lens(Tensor): Shape [B] + grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + + # query, key, value function + def qkv_fn(x): + q = self.norm_q(self.q(x)).view(b, s, n, d) + k = self.norm_k(self.k(x)).view(b, s, n, d) + v = self.v(x).view(b, s, n, d) + return q, k, v + + q, k, v = qkv_fn(x) + + x = flash_attention( + q=rope_apply(q, grid_sizes, freqs), + k=rope_apply(k, grid_sizes, freqs), + v=v, + k_lens=seq_lens, + window_size=self.window_size) + + # output + x = x.flatten(2) + x = self.o(x) + return x + + +class WanCrossAttention(WanSelfAttention): + + def forward(self, x, context, context_lens, crossattn_cache=None): + r""" + Args: + x(Tensor): Shape [B, L1, C] + context(Tensor): Shape [B, L2, C] + context_lens(Tensor): Shape [B] + crossattn_cache (List[dict], *optional*): Contains the cached key and value tensors for context embedding. + """ + b, n, d = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + q = self.norm_q(self.q(x)).view(b, -1, n, d) + + if crossattn_cache is not None: + if not crossattn_cache["is_init"]: + crossattn_cache["is_init"] = True + k = self.norm_k(self.k(context)).view(b, -1, n, d) + v = self.v(context).view(b, -1, n, d) + crossattn_cache["k"] = k + crossattn_cache["v"] = v + else: + k = crossattn_cache["k"] + v = crossattn_cache["v"] + else: + k = self.norm_k(self.k(context)).view(b, -1, n, d) + v = self.v(context).view(b, -1, n, d) + + # compute attention + x = flash_attention(q, k, v, k_lens=context_lens) + + # output + x = x.flatten(2) + x = self.o(x) + return x + + +class WanAttentionBlock(nn.Module): + + def __init__(self, + dim, + ffn_dim, + num_heads, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=False, + eps=1e-6): + super().__init__() + self.dim = dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + self.window_size = window_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + + # layers + self.norm1 = WanLayerNorm(dim, eps) + self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm, + eps) + self.norm3 = WanLayerNorm( + dim, eps, + elementwise_affine=True) if cross_attn_norm else nn.Identity() + self.cross_attn = WanCrossAttention(dim, num_heads, (-1, -1), qk_norm, + eps) + self.norm2 = WanLayerNorm(dim, eps) + self.ffn = nn.Sequential( + nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'), + nn.Linear(ffn_dim, dim)) + + # modulation + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + x, + e, + seq_lens, + grid_sizes, + freqs, + context, + context_lens, + ): + r""" + Args: + x(Tensor): Shape [B, L, C] + e(Tensor): Shape [B, L1, 6, C] + seq_lens(Tensor): Shape [B], length of each sequence in batch + grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + e = (self.modulation.unsqueeze(0) + e).chunk(6, dim=2) + + # self-attention + y = self.self_attn( + self.norm1(x).type_as(x) * (1 + e[1].squeeze(2)) + e[0].squeeze(2), + seq_lens, grid_sizes, freqs) + x = x + y * e[2].squeeze(2) + + # cross-attention & ffn function + def cross_attn_ffn(x, context, context_lens, e): + x = x + self.cross_attn(self.norm3(x), context, context_lens) + y = self.ffn( + self.norm2(x).type_as(x) * (1 + e[4].squeeze(2)) + e[3].squeeze(2)) + x = x + y * e[5].squeeze(2) + return x + + x = cross_attn_ffn(x, context, context_lens, e) + return x + + +class Head(nn.Module): + + def __init__(self, dim, out_dim, patch_size, eps=1e-6): + super().__init__() + self.dim = dim + self.out_dim = out_dim + self.patch_size = patch_size + self.eps = eps + + # layers + out_dim = math.prod(patch_size) * out_dim + self.norm = WanLayerNorm(dim, eps) + self.head = nn.Linear(dim, out_dim) + + # modulation + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x, e): + r""" + Args: + x(Tensor): Shape [B, L1, C] + e(Tensor): Shape [B, L1, C] + """ + e = (self.modulation.unsqueeze(0) + e.unsqueeze(2)).chunk(2, dim=2) + x = ( + self.head( + self.norm(x) * (1 + e[1].squeeze(2)) + e[0].squeeze(2))) + return x + + +class WanModel(ModelMixin, ConfigMixin): + r""" + Wan diffusion backbone supporting both text-to-video and image-to-video. + """ + + ignore_for_config = [ + 'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim', 'window_size' + ] + _no_split_modules = ['WanAttentionBlock'] + _supports_gradient_checkpointing = True + + @register_to_config + def __init__(self, + model_type='t2v', + patch_size=(1, 2, 2), + text_len=512, + in_dim=16, + dim=2048, + ffn_dim=8192, + freq_dim=256, + text_dim=4096, + out_dim=16, + num_heads=16, + num_layers=32, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=True, + eps=1e-6): + r""" + Initialize the diffusion model backbone. + + Args: + model_type (`str`, *optional*, defaults to 't2v'): + Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video) + patch_size (`tuple`, *optional*, defaults to (1, 2, 2)): + 3D patch dimensions for video embedding (t_patch, h_patch, w_patch) + text_len (`int`, *optional*, defaults to 512): + Fixed length for text embeddings + in_dim (`int`, *optional*, defaults to 16): + Input video channels (C_in) + dim (`int`, *optional*, defaults to 2048): + Hidden dimension of the transformer + ffn_dim (`int`, *optional*, defaults to 8192): + Intermediate dimension in feed-forward network + freq_dim (`int`, *optional*, defaults to 256): + Dimension for sinusoidal time embeddings + text_dim (`int`, *optional*, defaults to 4096): + Input dimension for text embeddings + out_dim (`int`, *optional*, defaults to 16): + Output video channels (C_out) + num_heads (`int`, *optional*, defaults to 16): + Number of attention heads + num_layers (`int`, *optional*, defaults to 32): + Number of transformer blocks + window_size (`tuple`, *optional*, defaults to (-1, -1)): + Window size for local attention (-1 indicates global attention) + qk_norm (`bool`, *optional*, defaults to True): + Enable query/key normalization + cross_attn_norm (`bool`, *optional*, defaults to False): + Enable cross-attention normalization + eps (`float`, *optional*, defaults to 1e-6): + Epsilon value for normalization layers + """ + + super().__init__() + + assert model_type in ['t2v', 'i2v', 'ti2v'] + self.model_type = model_type + + self.patch_size = patch_size + self.text_len = text_len + self.in_dim = in_dim + self.dim = dim + self.ffn_dim = ffn_dim + self.freq_dim = freq_dim + self.text_dim = text_dim + self.out_dim = out_dim + self.num_heads = num_heads + self.num_layers = num_layers + self.window_size = window_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + self.local_attn_size = 21 + + # embeddings + self.patch_embedding = nn.Conv3d( + in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'), + nn.Linear(dim, dim)) + + self.time_embedding = nn.Sequential( + nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6)) + + # blocks + self.blocks = nn.ModuleList([ + WanAttentionBlock(dim, ffn_dim, num_heads, window_size, qk_norm, + cross_attn_norm, eps) for _ in range(num_layers) + ]) + + # head + self.head = Head(dim, out_dim, patch_size, eps) + + # buffers (don't use register_buffer otherwise dtype will be changed in to()) + assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0 + d = dim // num_heads + self.freqs = torch.cat([ + rope_params(1024, d - 4 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + rope_params(1024, 2 * (d // 6)) + ], + dim=1) + + # initialize weights + self.init_weights() + + self.gradient_checkpointing = False + + def _set_gradient_checkpointing(self, module, value=False): + self.gradient_checkpointing = value + + def forward( + self, + *args, + **kwargs + ): + return self._forward(*args, **kwargs) + + def _forward( + self, + x, + t, + context, + seq_len, + y=None, + ): + r""" + Forward pass through the diffusion model + + Args: + x (List[Tensor]): + List of input video tensors, each with shape [C_in, F, H, W] + t (Tensor): + Diffusion timesteps tensor of shape [B] + context (List[Tensor]): + List of text embeddings each with shape [L, C] + seq_len (`int`): + Maximum sequence length for positional encoding + y (List[Tensor], *optional*): + Conditional video inputs for image-to-video mode, same shape as x + + Returns: + List[Tensor]: + List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8] + """ + if self.model_type == 'i2v': + assert y is not None + # params + device = self.patch_embedding.weight.device + if self.freqs.device != device: + self.freqs = self.freqs.to(device) + + if y is not None: + x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)] + + # embeddings + x = [self.patch_embedding(u.unsqueeze(0)) for u in x] + grid_sizes = torch.stack( + [torch.tensor(u.shape[2:], dtype=torch.long) for u in x]) + x = [u.flatten(2).transpose(1, 2) for u in x] + seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long) + assert seq_lens.max() <= seq_len + x = torch.cat([ + torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], + dim=1) for u in x + ]) + + # time embeddings + if t.dim() == 1: + t = t.unsqueeze(1).expand(-1, seq_len) + bt = t.size(0) + t = t.flatten() + e = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, + t).unflatten(0, (bt, seq_len)).type_as(x)) + e0 = self.time_projection(e).unflatten(2, (6, self.dim)) + + # context + context_lens = None + context = self.text_embedding( + torch.stack([ + torch.cat( + [u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) + for u in context + ])) + + # arguments + kwargs = dict( + e=e0, + seq_lens=seq_lens, + grid_sizes=grid_sizes, + freqs=self.freqs, + context=context, + context_lens=context_lens) + + def create_custom_forward(module): + def custom_forward(*inputs, **kwargs): + return module(*inputs, **kwargs) + return custom_forward + + for block in self.blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, **kwargs, + use_reentrant=False, + ) + else: + x = block(x, **kwargs) + + # head + x = self.head(x, e) + + # unpatchify + x = self.unpatchify(x, grid_sizes) + + return torch.stack(x) + + + def unpatchify(self, x, grid_sizes, c=None): + r""" + Reconstruct video tensors from patch embeddings. + + Args: + x (List[Tensor]): + List of patchified features, each with shape [L, C_out * prod(patch_size)] + grid_sizes (Tensor): + Original spatial-temporal grid dimensions before patching, + shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches) + + Returns: + List[Tensor]: + Reconstructed video tensors with shape [C_out, F, H / 8, W / 8] + """ + + c = self.out_dim if c is None else c + out = [] + for u, v in zip(x, grid_sizes.tolist()): + u = u[:math.prod(v)].view(*v, *self.patch_size, c) + u = torch.einsum('fhwpqrc->cfphqwr', u) + u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)]) + out.append(u) + return out + + def init_weights(self): + r""" + Initialize model parameters using Xavier initialization. + """ + + # basic init + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + if m.bias is not None: + nn.init.zeros_(m.bias) + + # init embeddings + nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1)) + for m in self.text_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=.02) + for m in self.time_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=.02) + + # init output layer + nn.init.zeros_(self.head.head.weight) diff --git a/wan_5b/modules/t5.py b/wan_5b/modules/t5.py new file mode 100644 index 0000000..6405e92 --- /dev/null +++ b/wan_5b/modules/t5.py @@ -0,0 +1,513 @@ +# Modified from transformers.models.t5.modeling_t5 +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import logging +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .tokenizers import HuggingfaceTokenizer + +__all__ = [ + 'T5Model', + 'T5Encoder', + 'T5Decoder', + 'T5EncoderModel', +] + + +def fp16_clamp(x): + if x.dtype == torch.float16 and torch.isinf(x).any(): + clamp = torch.finfo(x.dtype).max - 1000 + x = torch.clamp(x, min=-clamp, max=clamp) + return x + + +def init_weights(m): + if isinstance(m, T5LayerNorm): + nn.init.ones_(m.weight) + elif isinstance(m, T5Model): + nn.init.normal_(m.token_embedding.weight, std=1.0) + elif isinstance(m, T5FeedForward): + nn.init.normal_(m.gate[0].weight, std=m.dim**-0.5) + nn.init.normal_(m.fc1.weight, std=m.dim**-0.5) + nn.init.normal_(m.fc2.weight, std=m.dim_ffn**-0.5) + elif isinstance(m, T5Attention): + nn.init.normal_(m.q.weight, std=(m.dim * m.dim_attn)**-0.5) + nn.init.normal_(m.k.weight, std=m.dim**-0.5) + nn.init.normal_(m.v.weight, std=m.dim**-0.5) + nn.init.normal_(m.o.weight, std=(m.num_heads * m.dim_attn)**-0.5) + elif isinstance(m, T5RelativeEmbedding): + nn.init.normal_( + m.embedding.weight, std=(2 * m.num_buckets * m.num_heads)**-0.5) + + +class GELU(nn.Module): + + def forward(self, x): + return 0.5 * x * (1.0 + torch.tanh( + math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0)))) + + +class T5LayerNorm(nn.Module): + + def __init__(self, dim, eps=1e-6): + super(T5LayerNorm, self).__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + x = x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + + self.eps) + if self.weight.dtype in [torch.float16, torch.bfloat16]: + x = x.type_as(self.weight) + return self.weight * x + + +class T5Attention(nn.Module): + + def __init__(self, dim, dim_attn, num_heads, dropout=0.1): + assert dim_attn % num_heads == 0 + super(T5Attention, self).__init__() + self.dim = dim + self.dim_attn = dim_attn + self.num_heads = num_heads + self.head_dim = dim_attn // num_heads + + # layers + self.q = nn.Linear(dim, dim_attn, bias=False) + self.k = nn.Linear(dim, dim_attn, bias=False) + self.v = nn.Linear(dim, dim_attn, bias=False) + self.o = nn.Linear(dim_attn, dim, bias=False) + self.dropout = nn.Dropout(dropout) + + def forward(self, x, context=None, mask=None, pos_bias=None): + """ + x: [B, L1, C]. + context: [B, L2, C] or None. + mask: [B, L2] or [B, L1, L2] or None. + """ + # check inputs + context = x if context is None else context + b, n, c = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + q = self.q(x).view(b, -1, n, c) + k = self.k(context).view(b, -1, n, c) + v = self.v(context).view(b, -1, n, c) + + # attention bias + attn_bias = x.new_zeros(b, n, q.size(1), k.size(1)) + if pos_bias is not None: + attn_bias += pos_bias + if mask is not None: + assert mask.ndim in [2, 3] + mask = mask.view(b, 1, 1, + -1) if mask.ndim == 2 else mask.unsqueeze(1) + attn_bias.masked_fill_(mask == 0, torch.finfo(x.dtype).min) + + # compute attention (T5 does not use scaling) + attn = torch.einsum('binc,bjnc->bnij', q, k) + attn_bias + attn = F.softmax(attn.float(), dim=-1).type_as(attn) + x = torch.einsum('bnij,bjnc->binc', attn, v) + + # output + x = x.reshape(b, -1, n * c) + x = self.o(x) + x = self.dropout(x) + return x + + +class T5FeedForward(nn.Module): + + def __init__(self, dim, dim_ffn, dropout=0.1): + super(T5FeedForward, self).__init__() + self.dim = dim + self.dim_ffn = dim_ffn + + # layers + self.gate = nn.Sequential(nn.Linear(dim, dim_ffn, bias=False), GELU()) + self.fc1 = nn.Linear(dim, dim_ffn, bias=False) + self.fc2 = nn.Linear(dim_ffn, dim, bias=False) + self.dropout = nn.Dropout(dropout) + + def forward(self, x): + x = self.fc1(x) * self.gate(x) + x = self.dropout(x) + x = self.fc2(x) + x = self.dropout(x) + return x + + +class T5SelfAttention(nn.Module): + + def __init__(self, + dim, + dim_attn, + dim_ffn, + num_heads, + num_buckets, + shared_pos=True, + dropout=0.1): + super(T5SelfAttention, self).__init__() + self.dim = dim + self.dim_attn = dim_attn + self.dim_ffn = dim_ffn + self.num_heads = num_heads + self.num_buckets = num_buckets + self.shared_pos = shared_pos + + # layers + self.norm1 = T5LayerNorm(dim) + self.attn = T5Attention(dim, dim_attn, num_heads, dropout) + self.norm2 = T5LayerNorm(dim) + self.ffn = T5FeedForward(dim, dim_ffn, dropout) + self.pos_embedding = None if shared_pos else T5RelativeEmbedding( + num_buckets, num_heads, bidirectional=True) + + def forward(self, x, mask=None, pos_bias=None): + e = pos_bias if self.shared_pos else self.pos_embedding( + x.size(1), x.size(1)) + x = fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e)) + x = fp16_clamp(x + self.ffn(self.norm2(x))) + return x + + +class T5CrossAttention(nn.Module): + + def __init__(self, + dim, + dim_attn, + dim_ffn, + num_heads, + num_buckets, + shared_pos=True, + dropout=0.1): + super(T5CrossAttention, self).__init__() + self.dim = dim + self.dim_attn = dim_attn + self.dim_ffn = dim_ffn + self.num_heads = num_heads + self.num_buckets = num_buckets + self.shared_pos = shared_pos + + # layers + self.norm1 = T5LayerNorm(dim) + self.self_attn = T5Attention(dim, dim_attn, num_heads, dropout) + self.norm2 = T5LayerNorm(dim) + self.cross_attn = T5Attention(dim, dim_attn, num_heads, dropout) + self.norm3 = T5LayerNorm(dim) + self.ffn = T5FeedForward(dim, dim_ffn, dropout) + self.pos_embedding = None if shared_pos else T5RelativeEmbedding( + num_buckets, num_heads, bidirectional=False) + + def forward(self, + x, + mask=None, + encoder_states=None, + encoder_mask=None, + pos_bias=None): + e = pos_bias if self.shared_pos else self.pos_embedding( + x.size(1), x.size(1)) + x = fp16_clamp(x + self.self_attn(self.norm1(x), mask=mask, pos_bias=e)) + x = fp16_clamp(x + self.cross_attn( + self.norm2(x), context=encoder_states, mask=encoder_mask)) + x = fp16_clamp(x + self.ffn(self.norm3(x))) + return x + + +class T5RelativeEmbedding(nn.Module): + + def __init__(self, num_buckets, num_heads, bidirectional, max_dist=128): + super(T5RelativeEmbedding, self).__init__() + self.num_buckets = num_buckets + self.num_heads = num_heads + self.bidirectional = bidirectional + self.max_dist = max_dist + + # layers + self.embedding = nn.Embedding(num_buckets, num_heads) + + def forward(self, lq, lk): + device = self.embedding.weight.device + rel_pos = torch.arange(lk, device=device).unsqueeze(0) - \ + torch.arange(lq, device=device).unsqueeze(1) + rel_pos = self._relative_position_bucket(rel_pos) + rel_pos_embeds = self.embedding(rel_pos) + rel_pos_embeds = rel_pos_embeds.permute(2, 0, 1).unsqueeze( + 0) # [1, N, Lq, Lk] + return rel_pos_embeds.contiguous() + + def _relative_position_bucket(self, rel_pos): + # preprocess + if self.bidirectional: + num_buckets = self.num_buckets // 2 + rel_buckets = (rel_pos > 0).long() * num_buckets + rel_pos = torch.abs(rel_pos) + else: + num_buckets = self.num_buckets + rel_buckets = 0 + rel_pos = -torch.min(rel_pos, torch.zeros_like(rel_pos)) + + # embeddings for small and large positions + max_exact = num_buckets // 2 + rel_pos_large = max_exact + (torch.log(rel_pos.float() / max_exact) / + math.log(self.max_dist / max_exact) * + (num_buckets - max_exact)).long() + rel_pos_large = torch.min( + rel_pos_large, torch.full_like(rel_pos_large, num_buckets - 1)) + rel_buckets += torch.where(rel_pos < max_exact, rel_pos, rel_pos_large) + return rel_buckets + + +class T5Encoder(nn.Module): + + def __init__(self, + vocab, + dim, + dim_attn, + dim_ffn, + num_heads, + num_layers, + num_buckets, + shared_pos=True, + dropout=0.1): + super(T5Encoder, self).__init__() + self.dim = dim + self.dim_attn = dim_attn + self.dim_ffn = dim_ffn + self.num_heads = num_heads + self.num_layers = num_layers + self.num_buckets = num_buckets + self.shared_pos = shared_pos + + # layers + self.token_embedding = vocab if isinstance(vocab, nn.Embedding) \ + else nn.Embedding(vocab, dim) + self.pos_embedding = T5RelativeEmbedding( + num_buckets, num_heads, bidirectional=True) if shared_pos else None + self.dropout = nn.Dropout(dropout) + self.blocks = nn.ModuleList([ + T5SelfAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets, + shared_pos, dropout) for _ in range(num_layers) + ]) + self.norm = T5LayerNorm(dim) + + # initialize weights + self.apply(init_weights) + + def forward(self, ids, mask=None): + x = self.token_embedding(ids) + x = self.dropout(x) + e = self.pos_embedding(x.size(1), + x.size(1)) if self.shared_pos else None + for block in self.blocks: + x = block(x, mask, pos_bias=e) + x = self.norm(x) + x = self.dropout(x) + return x + + +class T5Decoder(nn.Module): + + def __init__(self, + vocab, + dim, + dim_attn, + dim_ffn, + num_heads, + num_layers, + num_buckets, + shared_pos=True, + dropout=0.1): + super(T5Decoder, self).__init__() + self.dim = dim + self.dim_attn = dim_attn + self.dim_ffn = dim_ffn + self.num_heads = num_heads + self.num_layers = num_layers + self.num_buckets = num_buckets + self.shared_pos = shared_pos + + # layers + self.token_embedding = vocab if isinstance(vocab, nn.Embedding) \ + else nn.Embedding(vocab, dim) + self.pos_embedding = T5RelativeEmbedding( + num_buckets, num_heads, bidirectional=False) if shared_pos else None + self.dropout = nn.Dropout(dropout) + self.blocks = nn.ModuleList([ + T5CrossAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets, + shared_pos, dropout) for _ in range(num_layers) + ]) + self.norm = T5LayerNorm(dim) + + # initialize weights + self.apply(init_weights) + + def forward(self, ids, mask=None, encoder_states=None, encoder_mask=None): + b, s = ids.size() + + # causal mask + if mask is None: + mask = torch.tril(torch.ones(1, s, s).to(ids.device)) + elif mask.ndim == 2: + mask = torch.tril(mask.unsqueeze(1).expand(-1, s, -1)) + + # layers + x = self.token_embedding(ids) + x = self.dropout(x) + e = self.pos_embedding(x.size(1), + x.size(1)) if self.shared_pos else None + for block in self.blocks: + x = block(x, mask, encoder_states, encoder_mask, pos_bias=e) + x = self.norm(x) + x = self.dropout(x) + return x + + +class T5Model(nn.Module): + + def __init__(self, + vocab_size, + dim, + dim_attn, + dim_ffn, + num_heads, + encoder_layers, + decoder_layers, + num_buckets, + shared_pos=True, + dropout=0.1): + super(T5Model, self).__init__() + self.vocab_size = vocab_size + self.dim = dim + self.dim_attn = dim_attn + self.dim_ffn = dim_ffn + self.num_heads = num_heads + self.encoder_layers = encoder_layers + self.decoder_layers = decoder_layers + self.num_buckets = num_buckets + + # layers + self.token_embedding = nn.Embedding(vocab_size, dim) + self.encoder = T5Encoder(self.token_embedding, dim, dim_attn, dim_ffn, + num_heads, encoder_layers, num_buckets, + shared_pos, dropout) + self.decoder = T5Decoder(self.token_embedding, dim, dim_attn, dim_ffn, + num_heads, decoder_layers, num_buckets, + shared_pos, dropout) + self.head = nn.Linear(dim, vocab_size, bias=False) + + # initialize weights + self.apply(init_weights) + + def forward(self, encoder_ids, encoder_mask, decoder_ids, decoder_mask): + x = self.encoder(encoder_ids, encoder_mask) + x = self.decoder(decoder_ids, decoder_mask, x, encoder_mask) + x = self.head(x) + return x + + +def _t5(name, + encoder_only=False, + decoder_only=False, + return_tokenizer=False, + tokenizer_kwargs={}, + dtype=torch.float32, + device='cpu', + **kwargs): + # sanity check + assert not (encoder_only and decoder_only) + + # params + if encoder_only: + model_cls = T5Encoder + kwargs['vocab'] = kwargs.pop('vocab_size') + kwargs['num_layers'] = kwargs.pop('encoder_layers') + _ = kwargs.pop('decoder_layers') + elif decoder_only: + model_cls = T5Decoder + kwargs['vocab'] = kwargs.pop('vocab_size') + kwargs['num_layers'] = kwargs.pop('decoder_layers') + _ = kwargs.pop('encoder_layers') + else: + model_cls = T5Model + + # init model + with torch.device(device): + model = model_cls(**kwargs) + + # set device + model = model.to(dtype=dtype, device=device) + + # init tokenizer + if return_tokenizer: + from .tokenizers import HuggingfaceTokenizer + tokenizer = HuggingfaceTokenizer(f'google/{name}', **tokenizer_kwargs) + return model, tokenizer + else: + return model + + +def umt5_xxl(**kwargs): + cfg = dict( + vocab_size=256384, + dim=4096, + dim_attn=4096, + dim_ffn=10240, + num_heads=64, + encoder_layers=24, + decoder_layers=24, + num_buckets=32, + shared_pos=False, + dropout=0.1) + cfg.update(**kwargs) + return _t5('umt5-xxl', **cfg) + + +class T5EncoderModel: + + def __init__( + self, + text_len, + dtype=torch.bfloat16, + device=None, + checkpoint_path=None, + tokenizer_path=None, + shard_fn=None, + ): + self.text_len = text_len + self.dtype = dtype + if device is None: + device = torch.cuda.current_device() if torch.cuda.is_available() else torch.device("cpu") + self.device = device + self.checkpoint_path = checkpoint_path + self.tokenizer_path = tokenizer_path + + # init model + model = umt5_xxl( + encoder_only=True, + return_tokenizer=False, + dtype=dtype, + device=device).eval().requires_grad_(False) + logging.info(f'loading {checkpoint_path}') + model.load_state_dict(torch.load(checkpoint_path, map_location='cpu')) + self.model = model + if shard_fn is not None: + self.model = shard_fn(self.model, sync_module_states=False) + else: + self.model.to(self.device) + # init tokenizer + self.tokenizer = HuggingfaceTokenizer( + name=tokenizer_path, seq_len=text_len, clean='whitespace') + + def __call__(self, texts, device): + ids, mask = self.tokenizer( + texts, return_mask=True, add_special_tokens=True) + ids = ids.to(device) + mask = mask.to(device) + seq_lens = mask.gt(0).sum(dim=1).long() + context = self.model(ids, mask) + return [u[:v] for u, v in zip(context, seq_lens)] diff --git a/wan_5b/modules/tokenizers.py b/wan_5b/modules/tokenizers.py new file mode 100644 index 0000000..121e591 --- /dev/null +++ b/wan_5b/modules/tokenizers.py @@ -0,0 +1,82 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import html +import string + +import ftfy +import regex as re +from transformers import AutoTokenizer + +__all__ = ['HuggingfaceTokenizer'] + + +def basic_clean(text): + text = ftfy.fix_text(text) + text = html.unescape(html.unescape(text)) + return text.strip() + + +def whitespace_clean(text): + text = re.sub(r'\s+', ' ', text) + text = text.strip() + return text + + +def canonicalize(text, keep_punctuation_exact_string=None): + text = text.replace('_', ' ') + if keep_punctuation_exact_string: + text = keep_punctuation_exact_string.join( + part.translate(str.maketrans('', '', string.punctuation)) + for part in text.split(keep_punctuation_exact_string)) + else: + text = text.translate(str.maketrans('', '', string.punctuation)) + text = text.lower() + text = re.sub(r'\s+', ' ', text) + return text.strip() + + +class HuggingfaceTokenizer: + + def __init__(self, name, seq_len=None, clean=None, **kwargs): + assert clean in (None, 'whitespace', 'lower', 'canonicalize') + self.name = name + self.seq_len = seq_len + self.clean = clean + + # init tokenizer + self.tokenizer = AutoTokenizer.from_pretrained(name, **kwargs) + self.vocab_size = self.tokenizer.vocab_size + + def __call__(self, sequence, **kwargs): + return_mask = kwargs.pop('return_mask', False) + + # arguments + _kwargs = {'return_tensors': 'pt'} + if self.seq_len is not None: + _kwargs.update({ + 'padding': 'max_length', + 'truncation': True, + 'max_length': self.seq_len + }) + _kwargs.update(**kwargs) + + # tokenization + if isinstance(sequence, str): + sequence = [sequence] + if self.clean: + sequence = [self._clean(u) for u in sequence] + ids = self.tokenizer(sequence, **_kwargs) + + # output + if return_mask: + return ids.input_ids, ids.attention_mask + else: + return ids.input_ids + + def _clean(self, text): + if self.clean == 'whitespace': + text = whitespace_clean(basic_clean(text)) + elif self.clean == 'lower': + text = whitespace_clean(basic_clean(text)).lower() + elif self.clean == 'canonicalize': + text = canonicalize(basic_clean(text)) + return text diff --git a/wan_5b/modules/vae2_1.py b/wan_5b/modules/vae2_1.py new file mode 100644 index 0000000..a5990fa --- /dev/null +++ b/wan_5b/modules/vae2_1.py @@ -0,0 +1,658 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import logging + +import torch +import torch.cuda.amp as amp +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +__all__ = [ + 'Wan2_1_VAE', +] + +CACHE_T = 2 + + +class CausalConv3d(nn.Conv3d): + """ + Causal 3d convolusion. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._padding = (self.padding[2], self.padding[2], self.padding[1], + self.padding[1], 2 * self.padding[0], 0) + self.padding = (0, 0, 0) + + def forward(self, x, cache_x=None): + padding = list(self._padding) + if cache_x is not None and self._padding[4] > 0: + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = F.pad(x, padding) + + return super().forward(x) + + +class RMS_norm(nn.Module): + + def __init__(self, dim, channel_first=True, images=True, bias=False): + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0. + + def forward(self, x): + return F.normalize( + x, dim=(1 if self.channel_first else + -1)) * self.scale * self.gamma + self.bias + + +class Upsample(nn.Upsample): + + def forward(self, x): + """ + Fix bfloat16 support for nearest neighbor interpolation. + """ + return super().forward(x.float()).type_as(x) + + +class Resample(nn.Module): + + def __init__(self, dim, mode): + assert mode in ('none', 'upsample2d', 'upsample3d', 'downsample2d', + 'downsample3d') + super().__init__() + self.dim = dim + self.mode = mode + + # layers + if mode == 'upsample2d': + self.resample = nn.Sequential( + Upsample(scale_factor=(2., 2.), mode='nearest-exact'), + nn.Conv2d(dim, dim // 2, 3, padding=1)) + elif mode == 'upsample3d': + self.resample = nn.Sequential( + Upsample(scale_factor=(2., 2.), mode='nearest-exact'), + nn.Conv2d(dim, dim // 2, 3, padding=1)) + self.time_conv = CausalConv3d( + dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) + + elif mode == 'downsample2d': + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2))) + elif mode == 'downsample3d': + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2))) + self.time_conv = CausalConv3d( + dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0)) + + else: + self.resample = nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + b, c, t, h, w = x.size() + if self.mode == 'upsample3d': + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = 'Rep' + feat_idx[0] += 1 + else: + + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[ + idx] is not None and feat_cache[idx] != 'Rep': + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + if cache_x.shape[2] < 2 and feat_cache[ + idx] is not None and feat_cache[idx] == 'Rep': + cache_x = torch.cat([ + torch.zeros_like(cache_x).to(cache_x.device), + cache_x + ], + dim=2) + if feat_cache[idx] == 'Rep': + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), + 3) + x = x.reshape(b, c, t * 2, h, w) + t = x.shape[2] + x = rearrange(x, 'b c t h w -> (b t) c h w') + x = self.resample(x) + x = rearrange(x, '(b t) c h w -> b c t h w', t=t) + + if self.mode == 'downsample3d': + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + + cache_x = x[:, :, -1:, :, :].clone() + + x = self.time_conv( + torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return x + + def init_weight(self, conv): + conv_weight = conv.weight + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + one_matrix = torch.eye(c1, c2) + init_matrix = one_matrix + nn.init.zeros_(conv_weight) + conv_weight.data[:, :, 1, 0, 0] = init_matrix + conv.weight.data.copy_(conv_weight) + nn.init.zeros_(conv.bias.data) + + def init_weight2(self, conv): + conv_weight = conv.weight.data + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + init_matrix = torch.eye(c1 // 2, c2) + conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix + conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix + conv.weight.data.copy_(conv_weight) + nn.init.zeros_(conv.bias.data) + + +class ResidualBlock(nn.Module): + + def __init__(self, in_dim, out_dim, dropout=0.0): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + # layers + self.residual = nn.Sequential( + RMS_norm(in_dim, images=False), nn.SiLU(), + CausalConv3d(in_dim, out_dim, 3, padding=1), + RMS_norm(out_dim, images=False), nn.SiLU(), nn.Dropout(dropout), + CausalConv3d(out_dim, out_dim, 3, padding=1)) + self.shortcut = CausalConv3d(in_dim, out_dim, 1) \ + if in_dim != out_dim else nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + h = self.shortcut(x) + for layer in self.residual: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + h + + +class AttentionBlock(nn.Module): + """ + Causal self-attention with a single head. + """ + + def __init__(self, dim): + super().__init__() + self.dim = dim + + # layers + self.norm = RMS_norm(dim) + self.to_qkv = nn.Conv2d(dim, dim * 3, 1) + self.proj = nn.Conv2d(dim, dim, 1) + + # zero out the last layer params + nn.init.zeros_(self.proj.weight) + + def forward(self, x): + identity = x + b, c, t, h, w = x.size() + x = rearrange(x, 'b c t h w -> (b t) c h w') + x = self.norm(x) + # compute query, key, value + q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, + -1).permute(0, 1, 3, + 2).contiguous().chunk( + 3, dim=-1) + + # apply attention + x = F.scaled_dot_product_attention( + q, + k, + v, + ) + x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) + + # output + x = self.proj(x) + x = rearrange(x, '(b t) c h w-> b c t h w', t=t) + return x + identity + + +class Encoder3d(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + + # dimensions + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + # init block + self.conv1 = CausalConv3d(3, dims[0], 3, padding=1) + + # downsample blocks + downsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + # residual (+attention) blocks + for _ in range(num_res_blocks): + downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + downsamples.append(AttentionBlock(out_dim)) + in_dim = out_dim + + # downsample block + if i != len(dim_mult) - 1: + mode = 'downsample3d' if temperal_downsample[ + i] else 'downsample2d' + downsamples.append(Resample(out_dim, mode=mode)) + scale /= 2.0 + self.downsamples = nn.Sequential(*downsamples) + + # middle blocks + self.middle = nn.Sequential( + ResidualBlock(out_dim, out_dim, dropout), AttentionBlock(out_dim), + ResidualBlock(out_dim, out_dim, dropout)) + + # output blocks + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), nn.SiLU(), + CausalConv3d(out_dim, z_dim, 3, padding=1)) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + ## downsamples + for layer in self.downsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## middle + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## head + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +class Decoder3d(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_upsample = temperal_upsample + + # dimensions + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + scale = 1.0 / 2**(len(dim_mult) - 2) + + # init block + self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) + + # middle blocks + self.middle = nn.Sequential( + ResidualBlock(dims[0], dims[0], dropout), AttentionBlock(dims[0]), + ResidualBlock(dims[0], dims[0], dropout)) + + # upsample blocks + upsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + # residual (+attention) blocks + if i == 1 or i == 2 or i == 3: + in_dim = in_dim // 2 + for _ in range(num_res_blocks + 1): + upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + upsamples.append(AttentionBlock(out_dim)) + in_dim = out_dim + + # upsample block + if i != len(dim_mult) - 1: + mode = 'upsample3d' if temperal_upsample[i] else 'upsample2d' + upsamples.append(Resample(out_dim, mode=mode)) + scale *= 2.0 + self.upsamples = nn.Sequential(*upsamples) + + # output blocks + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), nn.SiLU(), + CausalConv3d(out_dim, 3, 3, padding=1)) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + ## conv1 + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + ## middle + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## upsamples + for layer in self.upsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## head + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +def count_conv3d(model): + count = 0 + for m in model.modules(): + if isinstance(m, CausalConv3d): + count += 1 + return count + + +class WanVAE_(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + + # modules + self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks, + attn_scales, self.temperal_downsample, dropout) + self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.conv2 = CausalConv3d(z_dim, z_dim, 1) + self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks, + attn_scales, self.temperal_upsample, dropout) + + def forward(self, x): + mu, log_var = self.encode(x) + z = self.reparameterize(mu, log_var) + x_recon = self.decode(z) + return x_recon, mu, log_var + + def encode(self, x, scale): + self.clear_cache() + ## cache + t = x.shape[2] + iter_ = 1 + (t - 1) // 4 + # Split the encoder input over time as 1, 4, 4, 4, ... + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + out = self.encoder( + x[:, :, :1, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx) + else: + out_ = self.encoder( + x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx) + out = torch.cat([out, out_], 2) + mu, log_var = self.conv1(out).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view( + 1, self.z_dim, 1, 1, 1) + else: + mu = (mu - scale[0]) * scale[1] + self.clear_cache() + return mu + + def decode(self, z, scale): + self.clear_cache() + # z: [b,c,t,h,w] + if isinstance(scale[0], torch.Tensor): + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( + 1, self.z_dim, 1, 1, 1) + else: + z = z / scale[1] + scale[0] + iter_ = z.shape[2] + x = self.conv2(z) + for i in range(iter_): + self._conv_idx = [0] + if i == 0: + out = self.decoder( + x[:, :, i:i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx) + else: + out_ = self.decoder( + x[:, :, i:i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx) + out = torch.cat([out, out_], 2) + self.clear_cache() + return out + + def reparameterize(self, mu, log_var): + std = torch.exp(0.5 * log_var) + eps = torch.randn_like(std) + return eps * std + mu + + def sample(self, imgs, deterministic=False): + mu, log_var = self.encode(imgs) + if deterministic: + return mu + std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0)) + return mu + std * torch.randn_like(std) + + def clear_cache(self): + self._conv_num = count_conv3d(self.decoder) + self._conv_idx = [0] + self._feat_map = [None] * self._conv_num + #cache encode + self._enc_conv_num = count_conv3d(self.encoder) + self._enc_conv_idx = [0] + self._enc_feat_map = [None] * self._enc_conv_num + + +def _video_vae(pretrained_path=None, z_dim=None, device='cpu', **kwargs): + """ + Autoencoder3d adapted from Stable Diffusion 1.x, 2.x and XL. + """ + # params + cfg = dict( + dim=96, + z_dim=z_dim, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0) + cfg.update(**kwargs) + + # init model + with torch.device('meta'): + model = WanVAE_(**cfg) + + # load checkpoint + logging.info(f'loading {pretrained_path}') + model.load_state_dict( + torch.load(pretrained_path, map_location=device), assign=True) + + return model + + +class Wan2_1_VAE: + + def __init__(self, + z_dim=16, + vae_pth='cache/vae_step_411000.pth', + dtype=torch.float, + device="cuda"): + self.dtype = dtype + self.device = device + + mean = [ + -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, + 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921 + ] + std = [ + 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, + 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160 + ] + self.mean = torch.tensor(mean, dtype=dtype, device=device) + self.std = torch.tensor(std, dtype=dtype, device=device) + self.scale = [self.mean, 1.0 / self.std] + + # init model + self.model = _video_vae( + pretrained_path=vae_pth, + z_dim=z_dim, + ).eval().requires_grad_(False).to(device) + + def encode(self, videos): + """ + videos: A list of videos each with shape [C, T, H, W]. + """ + with amp.autocast(dtype=self.dtype): + return [ + self.model.encode(u.unsqueeze(0), self.scale).float().squeeze(0) + for u in videos + ] + + def decode(self, zs): + with amp.autocast(dtype=self.dtype): + return [ + self.model.decode(u.unsqueeze(0), + self.scale).float().clamp_(-1, 1).squeeze(0) + for u in zs + ] diff --git a/wan_5b/modules/vae2_2.py b/wan_5b/modules/vae2_2.py new file mode 100644 index 0000000..2237acc --- /dev/null +++ b/wan_5b/modules/vae2_2.py @@ -0,0 +1,1050 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import logging + +import torch +import torch.cuda.amp as amp +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +__all__ = [ + "Wan2_2_VAE", +] + +CACHE_T = 2 + + +class CausalConv3d(nn.Conv3d): + """ + Causal 3d convolusion. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._padding = ( + self.padding[2], + self.padding[2], + self.padding[1], + self.padding[1], + 2 * self.padding[0], + 0, + ) + self.padding = (0, 0, 0) + + def forward(self, x, cache_x=None): + padding = list(self._padding) + if cache_x is not None and self._padding[4] > 0: + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = F.pad(x, padding) + + return super().forward(x) + + +class RMS_norm(nn.Module): + + def __init__(self, dim, channel_first=True, images=True, bias=False): + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 + + def forward(self, x): + return (F.normalize(x, dim=(1 if self.channel_first else -1)) * + self.scale * self.gamma + self.bias) + + +class Upsample(nn.Upsample): + + def forward(self, x): + """ + Fix bfloat16 support for nearest neighbor interpolation. + """ + return super().forward(x.float()).type_as(x) + + +class Resample(nn.Module): + + def __init__(self, dim, mode): + assert mode in ( + "none", + "upsample2d", + "upsample3d", + "downsample2d", + "downsample3d", + ) + super().__init__() + self.dim = dim + self.mode = mode + + # layers + if mode == "upsample2d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim, 3, padding=1), + ) + elif mode == "upsample3d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim, 3, padding=1), + ) + self.time_conv = CausalConv3d( + dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) + elif mode == "downsample2d": + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2))) + elif mode == "downsample3d": + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2))) + self.time_conv = CausalConv3d( + dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0)) + else: + self.resample = nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + b, c, t, h, w = x.size() + if self.mode == "upsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = "Rep" + feat_idx[0] += 1 + else: + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if (cache_x.shape[2] < 2 and feat_cache[idx] is not None and + feat_cache[idx] != "Rep"): + # cache last frame of last two chunk + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), + cache_x, + ], + dim=2, + ) + if (cache_x.shape[2] < 2 and feat_cache[idx] is not None and + feat_cache[idx] == "Rep"): + cache_x = torch.cat( + [ + torch.zeros_like(cache_x).to(cache_x.device), + cache_x + ], + dim=2, + ) + if feat_cache[idx] == "Rep": + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), + 3) + x = x.reshape(b, c, t * 2, h, w) + t = x.shape[2] + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.resample(x) + x = rearrange(x, "(b t) c h w -> b c t h w", t=t) + + if self.mode == "downsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv( + torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return x + + def init_weight(self, conv): + conv_weight = conv.weight.detach().clone() + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + one_matrix = torch.eye(c1, c2) + init_matrix = one_matrix + nn.init.zeros_(conv_weight) + conv_weight.data[:, :, 1, 0, 0] = init_matrix + conv.weight = nn.Parameter(conv_weight) + nn.init.zeros_(conv.bias.data) + + def init_weight2(self, conv): + conv_weight = conv.weight.data.detach().clone() + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + init_matrix = torch.eye(c1 // 2, c2) + conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix + conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix + conv.weight = nn.Parameter(conv_weight) + nn.init.zeros_(conv.bias.data) + + +class ResidualBlock(nn.Module): + + def __init__(self, in_dim, out_dim, dropout=0.0): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + # layers + self.residual = nn.Sequential( + RMS_norm(in_dim, images=False), + nn.SiLU(), + CausalConv3d(in_dim, out_dim, 3, padding=1), + RMS_norm(out_dim, images=False), + nn.SiLU(), + nn.Dropout(dropout), + CausalConv3d(out_dim, out_dim, 3, padding=1), + ) + self.shortcut = ( + CausalConv3d(in_dim, out_dim, 1) + if in_dim != out_dim else nn.Identity()) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + h = self.shortcut(x) + for layer in self.residual: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), + cache_x, + ], + dim=2, + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + h + + +class AttentionBlock(nn.Module): + """ + Causal self-attention with a single head. + """ + + def __init__(self, dim): + super().__init__() + self.dim = dim + + # layers + self.norm = RMS_norm(dim) + self.to_qkv = nn.Conv2d(dim, dim * 3, 1) + self.proj = nn.Conv2d(dim, dim, 1) + + # zero out the last layer params + nn.init.zeros_(self.proj.weight) + + def forward(self, x): + identity = x + b, c, t, h, w = x.size() + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.norm(x) + # compute query, key, value + q, k, v = ( + self.to_qkv(x).reshape(b * t, 1, c * 3, + -1).permute(0, 1, 3, + 2).contiguous().chunk(3, dim=-1)) + + # apply attention + x = F.scaled_dot_product_attention( + q, + k, + v, + ) + x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) + + # output + x = self.proj(x) + x = rearrange(x, "(b t) c h w-> b c t h w", t=t) + return x + identity + + +def patchify(x, patch_size): + if patch_size == 1: + return x + if x.dim() == 4: + x = rearrange( + x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size) + elif x.dim() == 5: + x = rearrange( + x, + "b c f (h q) (w r) -> b (c r q) f h w", + q=patch_size, + r=patch_size, + ) + else: + raise ValueError(f"Invalid input shape: {x.shape}") + + return x + + +def unpatchify(x, patch_size): + if patch_size == 1: + return x + + if x.dim() == 4: + x = rearrange( + x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size) + elif x.dim() == 5: + x = rearrange( + x, + "b (c r q) f h w -> b c f (h q) (w r)", + q=patch_size, + r=patch_size, + ) + return x + + +class AvgDown3D(nn.Module): + + def __init__( + self, + in_channels, + out_channels, + factor_t, + factor_s=1, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + + assert in_channels * self.factor % out_channels == 0 + self.group_size = in_channels * self.factor // out_channels + + def forward(self, x: torch.Tensor) -> torch.Tensor: + pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t + pad = (0, 0, 0, 0, pad_t, 0) + x = F.pad(x, pad) + B, C, T, H, W = x.shape + x = x.view( + B, + C, + T // self.factor_t, + self.factor_t, + H // self.factor_s, + self.factor_s, + W // self.factor_s, + self.factor_s, + ) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view( + B, + C * self.factor, + T // self.factor_t, + H // self.factor_s, + W // self.factor_s, + ) + x = x.view( + B, + self.out_channels, + self.group_size, + T // self.factor_t, + H // self.factor_s, + W // self.factor_s, + ) + x = x.mean(dim=2) + return x + + +class DupUp3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + factor_t, + factor_s=1, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + + assert out_channels * self.factor % in_channels == 0 + self.repeats = out_channels * self.factor // in_channels + + def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor: + x = x.repeat_interleave(self.repeats, dim=1) + x = x.view( + x.size(0), + self.out_channels, + self.factor_t, + self.factor_s, + self.factor_s, + x.size(2), + x.size(3), + x.size(4), + ) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + x = x.view( + x.size(0), + self.out_channels, + x.size(2) * self.factor_t, + x.size(4) * self.factor_s, + x.size(6) * self.factor_s, + ) + if first_chunk: + x = x[:, :, self.factor_t - 1:, :, :] + return x + + +class Down_ResidualBlock(nn.Module): + + def __init__(self, + in_dim, + out_dim, + dropout, + mult, + temperal_downsample=False, + down_flag=False): + super().__init__() + + # Shortcut path with downsample + self.avg_shortcut = AvgDown3D( + in_dim, + out_dim, + factor_t=2 if temperal_downsample else 1, + factor_s=2 if down_flag else 1, + ) + + # Main path with residual blocks and downsample + downsamples = [] + for _ in range(mult): + downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + + # Add the final downsample block + if down_flag: + mode = "downsample3d" if temperal_downsample else "downsample2d" + downsamples.append(Resample(out_dim, mode=mode)) + + self.downsamples = nn.Sequential(*downsamples) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + x_copy = x.clone() + for module in self.downsamples: + x = module(x, feat_cache, feat_idx) + + return x + self.avg_shortcut(x_copy) + + +class Up_ResidualBlock(nn.Module): + + def __init__(self, + in_dim, + out_dim, + dropout, + mult, + temperal_upsample=False, + up_flag=False): + super().__init__() + # Shortcut path with upsample + if up_flag: + self.avg_shortcut = DupUp3D( + in_dim, + out_dim, + factor_t=2 if temperal_upsample else 1, + factor_s=2 if up_flag else 1, + ) + else: + self.avg_shortcut = None + + # Main path with residual blocks and upsample + upsamples = [] + for _ in range(mult): + upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + + # Add the final upsample block + if up_flag: + mode = "upsample3d" if temperal_upsample else "upsample2d" + upsamples.append(Resample(out_dim, mode=mode)) + + self.upsamples = nn.Sequential(*upsamples) + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + x_main = x.clone() + for module in self.upsamples: + x_main = module(x_main, feat_cache, feat_idx) + if self.avg_shortcut is not None: + x_shortcut = self.avg_shortcut(x, first_chunk) + return x_main + x_shortcut + else: + return x_main + + +class Encoder3d(nn.Module): + + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + + # dimensions + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + # init block + self.conv1 = CausalConv3d(12, dims[0], 3, padding=1) + + # downsample blocks + downsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + t_down_flag = ( + temperal_downsample[i] + if i < len(temperal_downsample) else False) + downsamples.append( + Down_ResidualBlock( + in_dim=in_dim, + out_dim=out_dim, + dropout=dropout, + mult=num_res_blocks, + temperal_downsample=t_down_flag, + down_flag=i != len(dim_mult) - 1, + )) + scale /= 2.0 + self.downsamples = nn.Sequential(*downsamples) + + # middle blocks + self.middle = nn.Sequential( + ResidualBlock(out_dim, out_dim, dropout), + AttentionBlock(out_dim), + ResidualBlock(out_dim, out_dim, dropout), + ) + + # # output blocks + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), + nn.SiLU(), + CausalConv3d(out_dim, z_dim, 3, padding=1), + ) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), + cache_x, + ], + dim=2, + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + ## downsamples + for layer in self.downsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## middle + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## head + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), + cache_x, + ], + dim=2, + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + + return x + + +class Decoder3d(nn.Module): + + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_upsample = temperal_upsample + + # dimensions + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + scale = 1.0 / 2**(len(dim_mult) - 2) + # init block + self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) + + # middle blocks + self.middle = nn.Sequential( + ResidualBlock(dims[0], dims[0], dropout), + AttentionBlock(dims[0]), + ResidualBlock(dims[0], dims[0], dropout), + ) + + # upsample blocks + upsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + t_up_flag = temperal_upsample[i] if i < len( + temperal_upsample) else False + upsamples.append( + Up_ResidualBlock( + in_dim=in_dim, + out_dim=out_dim, + dropout=dropout, + mult=num_res_blocks + 1, + temperal_upsample=t_up_flag, + up_flag=i != len(dim_mult) - 1, + )) + self.upsamples = nn.Sequential(*upsamples) + + # output blocks + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), + nn.SiLU(), + CausalConv3d(out_dim, 12, 3, padding=1), + ) + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), + cache_x, + ], + dim=2, + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## upsamples + for layer in self.upsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx, first_chunk) + else: + x = layer(x) + + ## head + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), + cache_x, + ], + dim=2, + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +def count_conv3d(model): + count = 0 + for m in model.modules(): + if isinstance(m, CausalConv3d): + count += 1 + return count + + +class WanVAE_(nn.Module): + + def __init__( + self, + dim=160, + dec_dim=256, + z_dim=16, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + + # modules + self.encoder = Encoder3d( + dim, + z_dim * 2, + dim_mult, + num_res_blocks, + attn_scales, + self.temperal_downsample, + dropout, + ) + self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.conv2 = CausalConv3d(z_dim, z_dim, 1) + self.decoder = Decoder3d( + dec_dim, + z_dim, + dim_mult, + num_res_blocks, + attn_scales, + self.temperal_upsample, + dropout, + ) + + def forward(self, x, scale=[0, 1]): + mu = self.encode(x, scale) + x_recon = self.decode(mu, scale) + return x_recon, mu + + def encode(self, x, scale): + self.clear_cache() + x = patchify(x, patch_size=2) + t = x.shape[2] + iter_ = 1 + (t - 1) // 4 + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + out = self.encoder( + x[:, :, :1, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) + else: + out_ = self.encoder( + x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) + out = torch.cat([out, out_], 2) + mu, log_var = self.conv1(out).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view( + 1, self.z_dim, 1, 1, 1) + else: + mu = (mu - scale[0]) * scale[1] + self.clear_cache() + return mu + + def decode(self, z, scale): + self.clear_cache() + if isinstance(scale[0], torch.Tensor): + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( + 1, self.z_dim, 1, 1, 1) + else: + z = z / scale[1] + scale[0] + iter_ = z.shape[2] + x = self.conv2(z) + for i in range(iter_): + self._conv_idx = [0] + if i == 0: + out = self.decoder( + x[:, :, i:i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + first_chunk=True, + ) + else: + out_ = self.decoder( + x[:, :, i:i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + ) + out = torch.cat([out, out_], 2) + out = unpatchify(out, patch_size=2) + self.clear_cache() + return out + + def reparameterize(self, mu, log_var): + std = torch.exp(0.5 * log_var) + eps = torch.randn_like(std) + return eps * std + mu + + def sample(self, imgs, deterministic=False): + mu, log_var = self.encode(imgs) + if deterministic: + return mu + std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0)) + return mu + std * torch.randn_like(std) + + def clear_cache(self): + self._conv_num = count_conv3d(self.decoder) + self._conv_idx = [0] + self._feat_map = [None] * self._conv_num + # cache encode + self._enc_conv_num = count_conv3d(self.encoder) + self._enc_conv_idx = [0] + self._enc_feat_map = [None] * self._enc_conv_num + + +def _video_vae(pretrained_path=None, z_dim=48, dim=160, device="cpu", **kwargs): + # params + cfg = dict( + dim=dim, + z_dim=z_dim, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0, + ) + cfg.update(**kwargs) + + # init model + with torch.device("meta"): + model = WanVAE_(**cfg) + + # load checkpoint + logging.info(f"loading {pretrained_path}") + model.load_state_dict( + torch.load(pretrained_path, map_location=device), assign=True) + + return model + + +class Wan2_2_VAE: + + def __init__( + self, + z_dim=48, + c_dim=160, + vae_pth=None, + dim_mult=[1, 2, 4, 4], + temperal_downsample=[False, True, True], + dtype=torch.float, + device="cuda", + ): + + self.dtype = dtype + self.device = device + + mean = torch.tensor( + [ + -0.2289, + -0.0052, + -0.1323, + -0.2339, + -0.2799, + 0.0174, + 0.1838, + 0.1557, + -0.1382, + 0.0542, + 0.2813, + 0.0891, + 0.1570, + -0.0098, + 0.0375, + -0.1825, + -0.2246, + -0.1207, + -0.0698, + 0.5109, + 0.2665, + -0.2108, + -0.2158, + 0.2502, + -0.2055, + -0.0322, + 0.1109, + 0.1567, + -0.0729, + 0.0899, + -0.2799, + -0.1230, + -0.0313, + -0.1649, + 0.0117, + 0.0723, + -0.2839, + -0.2083, + -0.0520, + 0.3748, + 0.0152, + 0.1957, + 0.1433, + -0.2944, + 0.3573, + -0.0548, + -0.1681, + -0.0667, + ], + dtype=dtype, + device=device, + ) + std = torch.tensor( + [ + 0.4765, + 1.0364, + 0.4514, + 1.1677, + 0.5313, + 0.4990, + 0.4818, + 0.5013, + 0.8158, + 1.0344, + 0.5894, + 1.0901, + 0.6885, + 0.6165, + 0.8454, + 0.4978, + 0.5759, + 0.3523, + 0.7135, + 0.6804, + 0.5833, + 1.4146, + 0.8986, + 0.5659, + 0.7069, + 0.5338, + 0.4889, + 0.4917, + 0.4069, + 0.4999, + 0.6866, + 0.4093, + 0.5709, + 0.6065, + 0.6415, + 0.4944, + 0.5726, + 1.2042, + 0.5458, + 1.6887, + 0.3971, + 1.0600, + 0.3943, + 0.5537, + 0.5444, + 0.4089, + 0.7468, + 0.7744, + ], + dtype=dtype, + device=device, + ) + self.scale = [mean, 1.0 / std] + + # init model + self.model = ( + _video_vae( + pretrained_path=vae_pth, + z_dim=z_dim, + dim=c_dim, + dim_mult=dim_mult, + temperal_downsample=temperal_downsample, + ).eval().requires_grad_(False).to(device)) + + def encode(self, videos): + try: + if not isinstance(videos, list): + raise TypeError("videos should be a list") + with amp.autocast(dtype=self.dtype): + return [ + self.model.encode(u.unsqueeze(0), + self.scale).float().squeeze(0) + for u in videos + ] + except TypeError as e: + logging.info(e) + return None + + def decode(self, zs): + try: + if not isinstance(zs, list): + raise TypeError("zs should be a list") + with amp.autocast(dtype=self.dtype): + return [ + self.model.decode(u.unsqueeze(0), + self.scale).float().clamp_(-1, + 1).squeeze(0) + for u in zs + ] + except TypeError as e: + logging.info(e) + return None diff --git a/wan_5b/text2video.py b/wan_5b/text2video.py new file mode 100644 index 0000000..7c79c66 --- /dev/null +++ b/wan_5b/text2video.py @@ -0,0 +1,378 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import gc +import logging +import math +import os +import random +import sys +import types +from contextlib import contextmanager +from functools import partial + +import torch +import torch.cuda.amp as amp +import torch.distributed as dist +from tqdm import tqdm + +from .distributed.fsdp import shard_model +from .distributed.sequence_parallel import sp_attn_forward, sp_dit_forward +from .distributed.util import get_world_size +from .modules.model import WanModel +from .modules.t5 import T5EncoderModel +from .modules.vae2_1 import Wan2_1_VAE +from .utils.fm_solvers import ( + FlowDPMSolverMultistepScheduler, + get_sampling_sigmas, + retrieve_timesteps, +) +from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler + + +class WanT2V: + + def __init__( + self, + config, + checkpoint_dir, + device_id=0, + rank=0, + t5_fsdp=False, + dit_fsdp=False, + use_sp=False, + t5_cpu=False, + init_on_cpu=True, + convert_model_dtype=False, + ): + r""" + Initializes the Wan text-to-video generation model components. + + Args: + config (EasyDict): + Object containing model parameters initialized from config.py + checkpoint_dir (`str`): + Path to directory containing model checkpoints + device_id (`int`, *optional*, defaults to 0): + Id of target GPU device + rank (`int`, *optional*, defaults to 0): + Process rank for distributed training + t5_fsdp (`bool`, *optional*, defaults to False): + Enable FSDP sharding for T5 model + dit_fsdp (`bool`, *optional*, defaults to False): + Enable FSDP sharding for DiT model + use_sp (`bool`, *optional*, defaults to False): + Enable distribution strategy of sequence parallel. + t5_cpu (`bool`, *optional*, defaults to False): + Whether to place T5 model on CPU. Only works without t5_fsdp. + init_on_cpu (`bool`, *optional*, defaults to True): + Enable initializing Transformer Model on CPU. Only works without FSDP or USP. + convert_model_dtype (`bool`, *optional*, defaults to False): + Convert DiT model parameters dtype to 'config.param_dtype'. + Only works without FSDP. + """ + self.device = torch.device(f"cuda:{device_id}") + self.config = config + self.rank = rank + self.t5_cpu = t5_cpu + self.init_on_cpu = init_on_cpu + + self.num_train_timesteps = config.num_train_timesteps + self.boundary = config.boundary + self.param_dtype = config.param_dtype + + if t5_fsdp or dit_fsdp or use_sp: + self.init_on_cpu = False + + shard_fn = partial(shard_model, device_id=device_id) + self.text_encoder = T5EncoderModel( + text_len=config.text_len, + dtype=config.t5_dtype, + device=torch.device('cpu'), + checkpoint_path=os.path.join(checkpoint_dir, config.t5_checkpoint), + tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer), + shard_fn=shard_fn if t5_fsdp else None) + + self.vae_stride = config.vae_stride + self.patch_size = config.patch_size + self.vae = Wan2_1_VAE( + vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint), + device=self.device) + + logging.info(f"Creating WanModel from {checkpoint_dir}") + self.low_noise_model = WanModel.from_pretrained( + checkpoint_dir, subfolder=config.low_noise_checkpoint) + self.low_noise_model = self._configure_model( + model=self.low_noise_model, + use_sp=use_sp, + dit_fsdp=dit_fsdp, + shard_fn=shard_fn, + convert_model_dtype=convert_model_dtype) + + self.high_noise_model = WanModel.from_pretrained( + checkpoint_dir, subfolder=config.high_noise_checkpoint) + self.high_noise_model = self._configure_model( + model=self.high_noise_model, + use_sp=use_sp, + dit_fsdp=dit_fsdp, + shard_fn=shard_fn, + convert_model_dtype=convert_model_dtype) + if use_sp: + self.sp_size = get_world_size() + else: + self.sp_size = 1 + + self.sample_neg_prompt = config.sample_neg_prompt + + def _configure_model(self, model, use_sp, dit_fsdp, shard_fn, + convert_model_dtype): + """ + Configures a model object. This includes setting evaluation modes, + applying distributed parallel strategy, and handling device placement. + + Args: + model (torch.nn.Module): + The model instance to configure. + use_sp (`bool`): + Enable distribution strategy of sequence parallel. + dit_fsdp (`bool`): + Enable FSDP sharding for DiT model. + shard_fn (callable): + The function to apply FSDP sharding. + convert_model_dtype (`bool`): + Convert DiT model parameters dtype to 'config.param_dtype'. + Only works without FSDP. + + Returns: + torch.nn.Module: + The configured model. + """ + model.eval().requires_grad_(False) + + if use_sp: + for block in model.blocks: + block.self_attn.forward = types.MethodType( + sp_attn_forward, block.self_attn) + model.forward = types.MethodType(sp_dit_forward, model) + + if dist.is_initialized(): + dist.barrier() + + if dit_fsdp: + model = shard_fn(model) + else: + if convert_model_dtype: + model.to(self.param_dtype) + if not self.init_on_cpu: + model.to(self.device) + + return model + + def _prepare_model_for_timestep(self, t, boundary, offload_model): + r""" + Prepares and returns the required model for the current timestep. + + Args: + t (torch.Tensor): + current timestep. + boundary (`int`): + The timestep threshold. If `t` is at or above this value, + the `high_noise_model` is considered as the required model. + offload_model (`bool`): + A flag intended to control the offloading behavior. + + Returns: + torch.nn.Module: + The active model on the target device for the current timestep. + """ + if t.item() >= boundary: + required_model_name = 'high_noise_model' + offload_model_name = 'low_noise_model' + else: + required_model_name = 'low_noise_model' + offload_model_name = 'high_noise_model' + if offload_model or self.init_on_cpu: + if next(getattr( + self, + offload_model_name).parameters()).device.type == 'cuda': + getattr(self, offload_model_name).to('cpu') + if next(getattr( + self, + required_model_name).parameters()).device.type == 'cpu': + getattr(self, required_model_name).to(self.device) + return getattr(self, required_model_name) + + def generate(self, + input_prompt, + size=(1280, 720), + frame_num=81, + shift=5.0, + sample_solver='unipc', + sampling_steps=50, + guide_scale=5.0, + n_prompt="", + seed=-1, + offload_model=True): + r""" + Generates video frames from text prompt using diffusion process. + + Args: + input_prompt (`str`): + Text prompt for content generation + size (`tuple[int]`, *optional*, defaults to (1280,720)): + Controls video resolution, (width,height). + frame_num (`int`, *optional*, defaults to 81): + How many frames to sample from a video. The number should be 4n+1 + shift (`float`, *optional*, defaults to 5.0): + Noise schedule shift parameter. Affects temporal dynamics + sample_solver (`str`, *optional*, defaults to 'unipc'): + Solver used to sample the video. + sampling_steps (`int`, *optional*, defaults to 50): + Number of diffusion sampling steps. Higher values improve quality but slow generation + guide_scale (`float` or tuple[`float`], *optional*, defaults 5.0): + Classifier-free guidance scale. Controls prompt adherence vs. creativity. + If tuple, the first guide_scale will be used for low noise model and + the second guide_scale will be used for high noise model. + n_prompt (`str`, *optional*, defaults to ""): + Negative prompt for content exclusion. If not given, use `config.sample_neg_prompt` + seed (`int`, *optional*, defaults to -1): + Random seed for noise generation. If -1, use random seed. + offload_model (`bool`, *optional*, defaults to True): + If True, offloads models to CPU during generation to save VRAM + + Returns: + torch.Tensor: + Generated video frames tensor. Dimensions: (C, N H, W) where: + - C: Color channels (3 for RGB) + - N: Number of frames (81) + - H: Frame height (from size) + - W: Frame width from size) + """ + # preprocess + guide_scale = (guide_scale, guide_scale) if isinstance( + guide_scale, float) else guide_scale + F = frame_num + target_shape = (self.vae.model.z_dim, (F - 1) // self.vae_stride[0] + 1, + size[1] // self.vae_stride[1], + size[0] // self.vae_stride[2]) + + seq_len = math.ceil((target_shape[2] * target_shape[3]) / + (self.patch_size[1] * self.patch_size[2]) * + target_shape[1] / self.sp_size) * self.sp_size + + if n_prompt == "": + n_prompt = self.sample_neg_prompt + seed = seed if seed >= 0 else random.randint(0, sys.maxsize) + seed_g = torch.Generator(device=self.device) + seed_g.manual_seed(seed) + + if not self.t5_cpu: + self.text_encoder.model.to(self.device) + context = self.text_encoder([input_prompt], self.device) + context_null = self.text_encoder([n_prompt], self.device) + if offload_model: + self.text_encoder.model.cpu() + else: + context = self.text_encoder([input_prompt], torch.device('cpu')) + context_null = self.text_encoder([n_prompt], torch.device('cpu')) + context = [t.to(self.device) for t in context] + context_null = [t.to(self.device) for t in context_null] + + noise = [ + torch.randn( + target_shape[0], + target_shape[1], + target_shape[2], + target_shape[3], + dtype=torch.float32, + device=self.device, + generator=seed_g) + ] + + @contextmanager + def noop_no_sync(): + yield + + no_sync_low_noise = getattr(self.low_noise_model, 'no_sync', + noop_no_sync) + no_sync_high_noise = getattr(self.high_noise_model, 'no_sync', + noop_no_sync) + + # evaluation mode + with ( + torch.amp.autocast('cuda', dtype=self.param_dtype), + torch.no_grad(), + no_sync_low_noise(), + no_sync_high_noise(), + ): + boundary = self.boundary * self.num_train_timesteps + + if sample_solver == 'unipc': + sample_scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=self.num_train_timesteps, + shift=1, + use_dynamic_shifting=False) + sample_scheduler.set_timesteps( + sampling_steps, device=self.device, shift=shift) + timesteps = sample_scheduler.timesteps + elif sample_solver == 'dpm++': + sample_scheduler = FlowDPMSolverMultistepScheduler( + num_train_timesteps=self.num_train_timesteps, + shift=1, + use_dynamic_shifting=False) + sampling_sigmas = get_sampling_sigmas(sampling_steps, shift) + timesteps, _ = retrieve_timesteps( + sample_scheduler, + device=self.device, + sigmas=sampling_sigmas) + else: + raise NotImplementedError("Unsupported solver.") + + # sample videos + latents = noise + + arg_c = {'context': context, 'seq_len': seq_len} + arg_null = {'context': context_null, 'seq_len': seq_len} + + for _, t in enumerate(tqdm(timesteps)): + latent_model_input = latents + timestep = [t] + + timestep = torch.stack(timestep) + + model = self._prepare_model_for_timestep( + t, boundary, offload_model) + sample_guide_scale = guide_scale[1] if t.item( + ) >= boundary else guide_scale[0] + + noise_pred_cond = model( + latent_model_input, t=timestep, **arg_c)[0] + noise_pred_uncond = model( + latent_model_input, t=timestep, **arg_null)[0] + + noise_pred = noise_pred_uncond + sample_guide_scale * ( + noise_pred_cond - noise_pred_uncond) + + temp_x0 = sample_scheduler.step( + noise_pred.unsqueeze(0), + t, + latents[0].unsqueeze(0), + return_dict=False, + generator=seed_g)[0] + latents = [temp_x0.squeeze(0)] + + x0 = latents + if offload_model: + self.low_noise_model.cpu() + self.high_noise_model.cpu() + torch.cuda.empty_cache() + if self.rank == 0: + videos = self.vae.decode(x0) + + del noise, latents + del sample_scheduler + if offload_model: + gc.collect() + torch.cuda.synchronize() + if dist.is_initialized(): + dist.barrier() + + return videos[0] if self.rank == 0 else None diff --git a/wan_5b/textimage2video.py b/wan_5b/textimage2video.py new file mode 100644 index 0000000..67e9fd2 --- /dev/null +++ b/wan_5b/textimage2video.py @@ -0,0 +1,619 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import gc +import logging +import math +import os +import random +import sys +import types +from contextlib import contextmanager +from functools import partial + +import torch +import torch.cuda.amp as amp +import torch.distributed as dist +import torchvision.transforms.functional as TF +from PIL import Image +from tqdm import tqdm + +from .distributed.fsdp import shard_model +from .distributed.sequence_parallel import sp_attn_forward, sp_dit_forward +from .distributed.util import get_world_size +from .modules.model import WanModel +from .modules.t5 import T5EncoderModel +from .modules.vae2_2 import Wan2_2_VAE +from .utils.fm_solvers import ( + FlowDPMSolverMultistepScheduler, + get_sampling_sigmas, + retrieve_timesteps, +) +from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler +from .utils.utils import best_output_size, masks_like + + +class WanTI2V: + + def __init__( + self, + config, + checkpoint_dir, + device_id=0, + rank=0, + t5_fsdp=False, + dit_fsdp=False, + use_sp=False, + t5_cpu=False, + init_on_cpu=True, + convert_model_dtype=False, + ): + r""" + Initializes the Wan text-to-video generation model components. + + Args: + config (EasyDict): + Object containing model parameters initialized from config.py + checkpoint_dir (`str`): + Path to directory containing model checkpoints + device_id (`int`, *optional*, defaults to 0): + Id of target GPU device + rank (`int`, *optional*, defaults to 0): + Process rank for distributed training + t5_fsdp (`bool`, *optional*, defaults to False): + Enable FSDP sharding for T5 model + dit_fsdp (`bool`, *optional*, defaults to False): + Enable FSDP sharding for DiT model + use_sp (`bool`, *optional*, defaults to False): + Enable distribution strategy of sequence parallel. + t5_cpu (`bool`, *optional*, defaults to False): + Whether to place T5 model on CPU. Only works without t5_fsdp. + init_on_cpu (`bool`, *optional*, defaults to True): + Enable initializing Transformer Model on CPU. Only works without FSDP or USP. + convert_model_dtype (`bool`, *optional*, defaults to False): + Convert DiT model parameters dtype to 'config.param_dtype'. + Only works without FSDP. + """ + self.device = torch.device(f"cuda:{device_id}") + self.config = config + self.rank = rank + self.t5_cpu = t5_cpu + self.init_on_cpu = init_on_cpu + + self.num_train_timesteps = config.num_train_timesteps + self.param_dtype = config.param_dtype + + if t5_fsdp or dit_fsdp or use_sp: + self.init_on_cpu = False + + shard_fn = partial(shard_model, device_id=device_id) + self.text_encoder = T5EncoderModel( + text_len=config.text_len, + dtype=config.t5_dtype, + device=torch.device('cpu'), + checkpoint_path=os.path.join(checkpoint_dir, config.t5_checkpoint), + tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer), + shard_fn=shard_fn if t5_fsdp else None) + + self.vae_stride = config.vae_stride + self.patch_size = config.patch_size + self.vae = Wan2_2_VAE( + vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint), + device=self.device) + + logging.info(f"Creating WanModel from {checkpoint_dir}") + self.model = WanModel.from_pretrained(checkpoint_dir) + self.model = self._configure_model( + model=self.model, + use_sp=use_sp, + dit_fsdp=dit_fsdp, + shard_fn=shard_fn, + convert_model_dtype=convert_model_dtype) + + if use_sp: + self.sp_size = get_world_size() + else: + self.sp_size = 1 + + self.sample_neg_prompt = config.sample_neg_prompt + + def _configure_model(self, model, use_sp, dit_fsdp, shard_fn, + convert_model_dtype): + """ + Configures a model object. This includes setting evaluation modes, + applying distributed parallel strategy, and handling device placement. + + Args: + model (torch.nn.Module): + The model instance to configure. + use_sp (`bool`): + Enable distribution strategy of sequence parallel. + dit_fsdp (`bool`): + Enable FSDP sharding for DiT model. + shard_fn (callable): + The function to apply FSDP sharding. + convert_model_dtype (`bool`): + Convert DiT model parameters dtype to 'config.param_dtype'. + Only works without FSDP. + + Returns: + torch.nn.Module: + The configured model. + """ + model.eval().requires_grad_(False) + + if use_sp: + for block in model.blocks: + block.self_attn.forward = types.MethodType( + sp_attn_forward, block.self_attn) + model.forward = types.MethodType(sp_dit_forward, model) + + if dist.is_initialized(): + dist.barrier() + + if dit_fsdp: + model = shard_fn(model) + else: + if convert_model_dtype: + model.to(self.param_dtype) + if not self.init_on_cpu: + model.to(self.device) + + return model + + def generate(self, + input_prompt, + img=None, + size=(1280, 704), + max_area=704 * 1280, + frame_num=81, + shift=5.0, + sample_solver='unipc', + sampling_steps=50, + guide_scale=5.0, + n_prompt="", + seed=-1, + offload_model=True): + r""" + Generates video frames from text prompt using diffusion process. + + Args: + input_prompt (`str`): + Text prompt for content generation + img (PIL.Image.Image): + Input image tensor. Shape: [3, H, W] + size (`tuple[int]`, *optional*, defaults to (1280,704)): + Controls video resolution, (width,height). + max_area (`int`, *optional*, defaults to 704*1280): + Maximum pixel area for latent space calculation. Controls video resolution scaling + frame_num (`int`, *optional*, defaults to 81): + How many frames to sample from a video. The number should be 4n+1 + shift (`float`, *optional*, defaults to 5.0): + Noise schedule shift parameter. Affects temporal dynamics + sample_solver (`str`, *optional*, defaults to 'unipc'): + Solver used to sample the video. + sampling_steps (`int`, *optional*, defaults to 50): + Number of diffusion sampling steps. Higher values improve quality but slow generation + guide_scale (`float`, *optional*, defaults 5.0): + Classifier-free guidance scale. Controls prompt adherence vs. creativity. + n_prompt (`str`, *optional*, defaults to ""): + Negative prompt for content exclusion. If not given, use `config.sample_neg_prompt` + seed (`int`, *optional*, defaults to -1): + Random seed for noise generation. If -1, use random seed. + offload_model (`bool`, *optional*, defaults to True): + If True, offloads models to CPU during generation to save VRAM + + Returns: + torch.Tensor: + Generated video frames tensor. Dimensions: (C, N H, W) where: + - C: Color channels (3 for RGB) + - N: Number of frames (81) + - H: Frame height (from size) + - W: Frame width from size) + """ + # i2v + if img is not None: + return self.i2v( + input_prompt=input_prompt, + img=img, + max_area=max_area, + frame_num=frame_num, + shift=shift, + sample_solver=sample_solver, + sampling_steps=sampling_steps, + guide_scale=guide_scale, + n_prompt=n_prompt, + seed=seed, + offload_model=offload_model) + # t2v + return self.t2v( + input_prompt=input_prompt, + size=size, + frame_num=frame_num, + shift=shift, + sample_solver=sample_solver, + sampling_steps=sampling_steps, + guide_scale=guide_scale, + n_prompt=n_prompt, + seed=seed, + offload_model=offload_model) + + def t2v(self, + input_prompt, + size=(1280, 704), + frame_num=121, + shift=5.0, + sample_solver='unipc', + sampling_steps=50, + guide_scale=5.0, + n_prompt="", + seed=-1, + offload_model=True): + r""" + Generates video frames from text prompt using diffusion process. + + Args: + input_prompt (`str`): + Text prompt for content generation + size (`tuple[int]`, *optional*, defaults to (1280,704)): + Controls video resolution, (width,height). + frame_num (`int`, *optional*, defaults to 121): + How many frames to sample from a video. The number should be 4n+1 + shift (`float`, *optional*, defaults to 5.0): + Noise schedule shift parameter. Affects temporal dynamics + sample_solver (`str`, *optional*, defaults to 'unipc'): + Solver used to sample the video. + sampling_steps (`int`, *optional*, defaults to 50): + Number of diffusion sampling steps. Higher values improve quality but slow generation + guide_scale (`float`, *optional*, defaults 5.0): + Classifier-free guidance scale. Controls prompt adherence vs. creativity. + n_prompt (`str`, *optional*, defaults to ""): + Negative prompt for content exclusion. If not given, use `config.sample_neg_prompt` + seed (`int`, *optional*, defaults to -1): + Random seed for noise generation. If -1, use random seed. + offload_model (`bool`, *optional*, defaults to True): + If True, offloads models to CPU during generation to save VRAM + + Returns: + torch.Tensor: + Generated video frames tensor. Dimensions: (C, N H, W) where: + - C: Color channels (3 for RGB) + - N: Number of frames (81) + - H: Frame height (from size) + - W: Frame width from size) + """ + # preprocess + F = frame_num + target_shape = (self.vae.model.z_dim, (F - 1) // self.vae_stride[0] + 1, + size[1] // self.vae_stride[1], + size[0] // self.vae_stride[2]) + + seq_len = math.ceil((target_shape[2] * target_shape[3]) / + (self.patch_size[1] * self.patch_size[2]) * + target_shape[1] / self.sp_size) * self.sp_size + + if n_prompt == "": + n_prompt = self.sample_neg_prompt + seed = seed if seed >= 0 else random.randint(0, sys.maxsize) + seed_g = torch.Generator(device=self.device) + seed_g.manual_seed(seed) + + if not self.t5_cpu: + self.text_encoder.model.to(self.device) + context = self.text_encoder([input_prompt], self.device) + context_null = self.text_encoder([n_prompt], self.device) + if offload_model: + self.text_encoder.model.cpu() + else: + context = self.text_encoder([input_prompt], torch.device('cpu')) + context_null = self.text_encoder([n_prompt], torch.device('cpu')) + context = [t.to(self.device) for t in context] + context_null = [t.to(self.device) for t in context_null] + + noise = [ + torch.randn( + target_shape[0], + target_shape[1], + target_shape[2], + target_shape[3], + dtype=torch.float32, + device=self.device, + generator=seed_g) + ] + + @contextmanager + def noop_no_sync(): + yield + + no_sync = getattr(self.model, 'no_sync', noop_no_sync) + + # evaluation mode + with ( + torch.amp.autocast('cuda', dtype=self.param_dtype), + torch.no_grad(), + no_sync(), + ): + + if sample_solver == 'unipc': + sample_scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=self.num_train_timesteps, + shift=1, + use_dynamic_shifting=False) + sample_scheduler.set_timesteps( + sampling_steps, device=self.device, shift=shift) + timesteps = sample_scheduler.timesteps + elif sample_solver == 'dpm++': + sample_scheduler = FlowDPMSolverMultistepScheduler( + num_train_timesteps=self.num_train_timesteps, + shift=1, + use_dynamic_shifting=False) + sampling_sigmas = get_sampling_sigmas(sampling_steps, shift) + timesteps, _ = retrieve_timesteps( + sample_scheduler, + device=self.device, + sigmas=sampling_sigmas) + else: + raise NotImplementedError("Unsupported solver.") + + # sample videos + latents = noise + mask1, mask2 = masks_like(noise, zero=False) + + arg_c = {'context': context, 'seq_len': seq_len} + arg_null = {'context': context_null, 'seq_len': seq_len} + + if offload_model or self.init_on_cpu: + self.model.to(self.device) + torch.cuda.empty_cache() + + for _, t in enumerate(tqdm(timesteps)): + latent_model_input = latents + timestep = [t] + + timestep = torch.stack(timestep) + + temp_ts = (mask2[0][0][:, ::2, ::2] * timestep).flatten() + temp_ts = torch.cat([ + temp_ts, + temp_ts.new_ones(seq_len - temp_ts.size(0)) * timestep + ]) + timestep = temp_ts.unsqueeze(0) + + noise_pred_cond = self.model( + latent_model_input, t=timestep, **arg_c)[0] + noise_pred_uncond = self.model( + latent_model_input, t=timestep, **arg_null)[0] + + noise_pred = noise_pred_uncond + guide_scale * ( + noise_pred_cond - noise_pred_uncond) + + temp_x0 = sample_scheduler.step( + noise_pred.unsqueeze(0), + t, + latents[0].unsqueeze(0), + return_dict=False, + generator=seed_g)[0] + latents = [temp_x0.squeeze(0)] + x0 = latents + if offload_model: + self.model.cpu() + torch.cuda.synchronize() + torch.cuda.empty_cache() + if self.rank == 0: + videos = self.vae.decode(x0) + + del noise, latents + del sample_scheduler + if offload_model: + gc.collect() + torch.cuda.synchronize() + if dist.is_initialized(): + dist.barrier() + + return videos[0] if self.rank == 0 else None + + def i2v(self, + input_prompt, + img, + max_area=704 * 1280, + frame_num=121, + shift=5.0, + sample_solver='unipc', + sampling_steps=40, + guide_scale=5.0, + n_prompt="", + seed=-1, + offload_model=True): + r""" + Generates video frames from input image and text prompt using diffusion process. + + Args: + input_prompt (`str`): + Text prompt for content generation. + img (PIL.Image.Image): + Input image tensor. Shape: [3, H, W] + max_area (`int`, *optional*, defaults to 704*1280): + Maximum pixel area for latent space calculation. Controls video resolution scaling + frame_num (`int`, *optional*, defaults to 121): + How many frames to sample from a video. The number should be 4n+1 + shift (`float`, *optional*, defaults to 5.0): + Noise schedule shift parameter. Affects temporal dynamics + [NOTE]: If you want to generate a 480p video, it is recommended to set the shift value to 3.0. + sample_solver (`str`, *optional*, defaults to 'unipc'): + Solver used to sample the video. + sampling_steps (`int`, *optional*, defaults to 40): + Number of diffusion sampling steps. Higher values improve quality but slow generation + guide_scale (`float`, *optional*, defaults 5.0): + Classifier-free guidance scale. Controls prompt adherence vs. creativity. + n_prompt (`str`, *optional*, defaults to ""): + Negative prompt for content exclusion. If not given, use `config.sample_neg_prompt` + seed (`int`, *optional*, defaults to -1): + Random seed for noise generation. If -1, use random seed + offload_model (`bool`, *optional*, defaults to True): + If True, offloads models to CPU during generation to save VRAM + + Returns: + torch.Tensor: + Generated video frames tensor. Dimensions: (C, N H, W) where: + - C: Color channels (3 for RGB) + - N: Number of frames (121) + - H: Frame height (from max_area) + - W: Frame width (from max_area) + """ + # preprocess + ih, iw = img.height, img.width + dh, dw = self.patch_size[1] * self.vae_stride[1], self.patch_size[ + 2] * self.vae_stride[2] + ow, oh = best_output_size(iw, ih, dw, dh, max_area) + + scale = max(ow / iw, oh / ih) + img = img.resize((round(iw * scale), round(ih * scale)), Image.LANCZOS) + + # center-crop + x1 = (img.width - ow) // 2 + y1 = (img.height - oh) // 2 + img = img.crop((x1, y1, x1 + ow, y1 + oh)) + assert img.width == ow and img.height == oh + + # to tensor + img = TF.to_tensor(img).sub_(0.5).div_(0.5).to(self.device).unsqueeze(1) + + F = frame_num + seq_len = ((F - 1) // self.vae_stride[0] + 1) * ( + oh // self.vae_stride[1]) * (ow // self.vae_stride[2]) // ( + self.patch_size[1] * self.patch_size[2]) + seq_len = int(math.ceil(seq_len / self.sp_size)) * self.sp_size + + seed = seed if seed >= 0 else random.randint(0, sys.maxsize) + seed_g = torch.Generator(device=self.device) + seed_g.manual_seed(seed) + noise = torch.randn( + self.vae.model.z_dim, (F - 1) // self.vae_stride[0] + 1, + oh // self.vae_stride[1], + ow // self.vae_stride[2], + dtype=torch.float32, + generator=seed_g, + device=self.device) + + if n_prompt == "": + n_prompt = self.sample_neg_prompt + + # preprocess + if not self.t5_cpu: + self.text_encoder.model.to(self.device) + context = self.text_encoder([input_prompt], self.device) + context_null = self.text_encoder([n_prompt], self.device) + if offload_model: + self.text_encoder.model.cpu() + else: + context = self.text_encoder([input_prompt], torch.device('cpu')) + context_null = self.text_encoder([n_prompt], torch.device('cpu')) + context = [t.to(self.device) for t in context] + context_null = [t.to(self.device) for t in context_null] + + z = self.vae.encode([img]) + + @contextmanager + def noop_no_sync(): + yield + + no_sync = getattr(self.model, 'no_sync', noop_no_sync) + + # evaluation mode + with ( + torch.amp.autocast('cuda', dtype=self.param_dtype), + torch.no_grad(), + no_sync(), + ): + + if sample_solver == 'unipc': + sample_scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=self.num_train_timesteps, + shift=1, + use_dynamic_shifting=False) + sample_scheduler.set_timesteps( + sampling_steps, device=self.device, shift=shift) + timesteps = sample_scheduler.timesteps + elif sample_solver == 'dpm++': + sample_scheduler = FlowDPMSolverMultistepScheduler( + num_train_timesteps=self.num_train_timesteps, + shift=1, + use_dynamic_shifting=False) + sampling_sigmas = get_sampling_sigmas(sampling_steps, shift) + timesteps, _ = retrieve_timesteps( + sample_scheduler, + device=self.device, + sigmas=sampling_sigmas) + else: + raise NotImplementedError("Unsupported solver.") + + # sample videos + latent = noise + mask1, mask2 = masks_like([noise], zero=True) + latent = (1. - mask2[0]) * z[0] + mask2[0] * latent + + arg_c = { + 'context': [context[0]], + 'seq_len': seq_len, + } + + arg_null = { + 'context': context_null, + 'seq_len': seq_len, + } + + if offload_model or self.init_on_cpu: + self.model.to(self.device) + torch.cuda.empty_cache() + + for _, t in enumerate(tqdm(timesteps)): + latent_model_input = [latent.to(self.device)] + timestep = [t] + + timestep = torch.stack(timestep).to(self.device) + + temp_ts = (mask2[0][0][:, ::2, ::2] * timestep).flatten() + temp_ts = torch.cat([ + temp_ts, + temp_ts.new_ones(seq_len - temp_ts.size(0)) * timestep + ]) + timestep = temp_ts.unsqueeze(0) + + noise_pred_cond = self.model( + latent_model_input, t=timestep, **arg_c)[0] + if offload_model: + torch.cuda.empty_cache() + noise_pred_uncond = self.model( + latent_model_input, t=timestep, **arg_null)[0] + if offload_model: + torch.cuda.empty_cache() + noise_pred = noise_pred_uncond + guide_scale * ( + noise_pred_cond - noise_pred_uncond) + + temp_x0 = sample_scheduler.step( + noise_pred.unsqueeze(0), + t, + latent.unsqueeze(0), + return_dict=False, + generator=seed_g)[0] + latent = temp_x0.squeeze(0) + latent = (1. - mask2[0]) * z[0] + mask2[0] * latent + + x0 = [latent] + del latent_model_input, timestep + + if offload_model: + self.model.cpu() + torch.cuda.synchronize() + torch.cuda.empty_cache() + + if self.rank == 0: + videos = self.vae.decode(x0) + + del noise, latent, x0 + del sample_scheduler + if offload_model: + gc.collect() + torch.cuda.synchronize() + if dist.is_initialized(): + dist.barrier() + + return videos[0] if self.rank == 0 else None diff --git a/wan_5b/utils/__init__.py b/wan_5b/utils/__init__.py new file mode 100644 index 0000000..5b17310 --- /dev/null +++ b/wan_5b/utils/__init__.py @@ -0,0 +1,12 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +from .fm_solvers import ( + FlowDPMSolverMultistepScheduler, + get_sampling_sigmas, + retrieve_timesteps, +) +from .fm_solvers_unipc import FlowUniPCMultistepScheduler + +__all__ = [ + 'HuggingfaceTokenizer', 'get_sampling_sigmas', 'retrieve_timesteps', + 'FlowDPMSolverMultistepScheduler', 'FlowUniPCMultistepScheduler' +] diff --git a/wan_5b/utils/fm_solvers.py b/wan_5b/utils/fm_solvers.py new file mode 100644 index 0000000..206a56b --- /dev/null +++ b/wan_5b/utils/fm_solvers.py @@ -0,0 +1,854 @@ +# Copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py +# Convert dpm solver for flow matching +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. + +import inspect +import math +from typing import List, Optional, Tuple, Union + +import numpy as np +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import ( + KarrasDiffusionSchedulers, + SchedulerMixin, + SchedulerOutput, +) +from diffusers.utils import deprecate, is_scipy_available +from diffusers.utils.torch_utils import randn_tensor + +if is_scipy_available(): + pass + + +def get_sampling_sigmas(sampling_steps, shift): + sigma = np.linspace(1, 0, sampling_steps + 1)[:sampling_steps] + sigma = (shift * sigma / (1 + (shift - 1) * sigma)) + + return sigma + + +def retrieve_timesteps( + scheduler, + num_inference_steps=None, + device=None, + timesteps=None, + sigmas=None, + **kwargs, +): + if timesteps is not None and sigmas is not None: + raise ValueError( + "Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values" + ) + if timesteps is not None: + accepts_timesteps = "timesteps" in set( + inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set( + inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class FlowDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin): + """ + `FlowDPMSolverMultistepScheduler` is a fast dedicated high-order solver for diffusion ODEs. + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps to train the model. This determines the resolution of the diffusion process. + solver_order (`int`, defaults to 2): + The DPMSolver order which can be `1`, `2`, or `3`. It is recommended to use `solver_order=2` for guided + sampling, and `solver_order=3` for unconditional sampling. This affects the number of model outputs stored + and used in multistep updates. + prediction_type (`str`, defaults to "flow_prediction"): + Prediction type of the scheduler function; must be `flow_prediction` for this scheduler, which predicts + the flow of the diffusion process. + shift (`float`, *optional*, defaults to 1.0): + A factor used to adjust the sigmas in the noise schedule. It modifies the step sizes during the sampling + process. + use_dynamic_shifting (`bool`, defaults to `False`): + Whether to apply dynamic shifting to the timesteps based on image resolution. If `True`, the shifting is + applied on the fly. + thresholding (`bool`, defaults to `False`): + Whether to use the "dynamic thresholding" method. This method adjusts the predicted sample to prevent + saturation and improve photorealism. + dynamic_thresholding_ratio (`float`, defaults to 0.995): + The ratio for the dynamic thresholding method. Valid only when `thresholding=True`. + sample_max_value (`float`, defaults to 1.0): + The threshold value for dynamic thresholding. Valid only when `thresholding=True` and + `algorithm_type="dpmsolver++"`. + algorithm_type (`str`, defaults to `dpmsolver++`): + Algorithm type for the solver; can be `dpmsolver`, `dpmsolver++`, `sde-dpmsolver` or `sde-dpmsolver++`. The + `dpmsolver` type implements the algorithms in the [DPMSolver](https://huggingface.co/papers/2206.00927) + paper, and the `dpmsolver++` type implements the algorithms in the + [DPMSolver++](https://huggingface.co/papers/2211.01095) paper. It is recommended to use `dpmsolver++` or + `sde-dpmsolver++` with `solver_order=2` for guided sampling like in Stable Diffusion. + solver_type (`str`, defaults to `midpoint`): + Solver type for the second-order solver; can be `midpoint` or `heun`. The solver type slightly affects the + sample quality, especially for a small number of steps. It is recommended to use `midpoint` solvers. + lower_order_final (`bool`, defaults to `True`): + Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can + stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10. + euler_at_final (`bool`, defaults to `False`): + Whether to use Euler's method in the final step. It is a trade-off between numerical stability and detail + richness. This can stabilize the sampling of the SDE variant of DPMSolver for small number of inference + steps, but sometimes may result in blurring. + final_sigmas_type (`str`, *optional*, defaults to "zero"): + The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final + sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0. + lambda_min_clipped (`float`, defaults to `-inf`): + Clipping threshold for the minimum value of `lambda(t)` for numerical stability. This is critical for the + cosine (`squaredcos_cap_v2`) noise schedule. + variance_type (`str`, *optional*): + Set to "learned" or "learned_range" for diffusion models that predict variance. If set, the model's output + contains the predicted Gaussian variance. + """ + + _compatibles = [e.name for e in KarrasDiffusionSchedulers] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + solver_order: int = 2, + prediction_type: str = "flow_prediction", + shift: Optional[float] = 1.0, + use_dynamic_shifting=False, + thresholding: bool = False, + dynamic_thresholding_ratio: float = 0.995, + sample_max_value: float = 1.0, + algorithm_type: str = "dpmsolver++", + solver_type: str = "midpoint", + lower_order_final: bool = True, + euler_at_final: bool = False, + final_sigmas_type: Optional[str] = "zero", # "zero", "sigma_min" + lambda_min_clipped: float = -float("inf"), + variance_type: Optional[str] = None, + invert_sigmas: bool = False, + ): + if algorithm_type in ["dpmsolver", "sde-dpmsolver"]: + deprecation_message = f"algorithm_type {algorithm_type} is deprecated and will be removed in a future version. Choose from `dpmsolver++` or `sde-dpmsolver++` instead" + deprecate("algorithm_types dpmsolver and sde-dpmsolver", "1.0.0", + deprecation_message) + + # settings for DPM-Solver + if algorithm_type not in [ + "dpmsolver", "dpmsolver++", "sde-dpmsolver", "sde-dpmsolver++" + ]: + if algorithm_type == "deis": + self.register_to_config(algorithm_type="dpmsolver++") + else: + raise NotImplementedError( + f"{algorithm_type} is not implemented for {self.__class__}") + + if solver_type not in ["midpoint", "heun"]: + if solver_type in ["logrho", "bh1", "bh2"]: + self.register_to_config(solver_type="midpoint") + else: + raise NotImplementedError( + f"{solver_type} is not implemented for {self.__class__}") + + if algorithm_type not in ["dpmsolver++", "sde-dpmsolver++" + ] and final_sigmas_type == "zero": + raise ValueError( + f"`final_sigmas_type` {final_sigmas_type} is not supported for `algorithm_type` {algorithm_type}. Please choose `sigma_min` instead." + ) + + # setable values + self.num_inference_steps = None + alphas = np.linspace(1, 1 / num_train_timesteps, + num_train_timesteps)[::-1].copy() + sigmas = 1.0 - alphas + sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32) + + if not use_dynamic_shifting: + # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution + sigmas = shift * sigmas / (1 + + (shift - 1) * sigmas) # pyright: ignore + + self.sigmas = sigmas + self.timesteps = sigmas * num_train_timesteps + + self.model_outputs = [None] * solver_order + self.lower_order_nums = 0 + self._step_index = None + self._begin_index = None + + self.sigma_min = self.sigmas[-1].item() + self.sigma_max = self.sigmas[0].item() + + @property + def step_index(self): + """ + The index counter for current timestep. It will increase 1 after each scheduler step. + """ + return self._step_index + + @property + def begin_index(self): + """ + The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + """ + return self._begin_index + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index + def set_begin_index(self, begin_index: int = 0): + """ + Sets the begin index for the scheduler. This function should be run from pipeline before the inference. + Args: + begin_index (`int`): + The begin index for the scheduler. + """ + self._begin_index = begin_index + + # Modified from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.set_timesteps + def set_timesteps( + self, + num_inference_steps: Union[int, None] = None, + device: Union[str, torch.device] = None, + sigmas: Optional[List[float]] = None, + mu: Optional[Union[float, None]] = None, + shift: Optional[Union[float, None]] = None, + ): + """ + Sets the discrete timesteps used for the diffusion chain (to be run before inference). + Args: + num_inference_steps (`int`): + Total number of the spacing of the time steps. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + """ + + if self.config.use_dynamic_shifting and mu is None: + raise ValueError( + " you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`" + ) + + if sigmas is None: + sigmas = np.linspace(self.sigma_max, self.sigma_min, + num_inference_steps + + 1).copy()[:-1] # pyright: ignore + + if self.config.use_dynamic_shifting: + sigmas = self.time_shift(mu, 1.0, sigmas) # pyright: ignore + else: + if shift is None: + shift = self.config.shift + sigmas = shift * sigmas / (1 + + (shift - 1) * sigmas) # pyright: ignore + + if self.config.final_sigmas_type == "sigma_min": + sigma_last = ((1 - self.alphas_cumprod[0]) / + self.alphas_cumprod[0])**0.5 + elif self.config.final_sigmas_type == "zero": + sigma_last = 0 + else: + raise ValueError( + f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" + ) + + timesteps = sigmas * self.config.num_train_timesteps + sigmas = np.concatenate([sigmas, [sigma_last] + ]).astype(np.float32) # pyright: ignore + + self.sigmas = torch.from_numpy(sigmas) + self.timesteps = torch.from_numpy(timesteps).to( + device=device, dtype=torch.int64) + + self.num_inference_steps = len(timesteps) + + self.model_outputs = [ + None, + ] * self.config.solver_order + self.lower_order_nums = 0 + + self._step_index = None + self._begin_index = None + # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample + def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: + """ + "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the + prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by + s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing + pixels from saturation at each step. We find that dynamic thresholding results in significantly better + photorealism as well as better image-text alignment, especially when using very large guidance weights." + https://arxiv.org/abs/2205.11487 + """ + dtype = sample.dtype + batch_size, channels, *remaining_dims = sample.shape + + if dtype not in (torch.float32, torch.float64): + sample = sample.float( + ) # upcast for quantile calculation, and clamp not implemented for cpu half + + # Flatten sample for doing quantile calculation along each image + sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) + + abs_sample = sample.abs() # "a certain percentile absolute pixel value" + + s = torch.quantile( + abs_sample, self.config.dynamic_thresholding_ratio, dim=1) + s = torch.clamp( + s, min=1, max=self.config.sample_max_value + ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] + s = s.unsqueeze( + 1) # (batch_size, 1) because clamp will broadcast along dim=0 + sample = torch.clamp( + sample, -s, s + ) / s # "we threshold xt0 to the range [-s, s] and then divide by s" + + sample = sample.reshape(batch_size, channels, *remaining_dims) + sample = sample.to(dtype) + + return sample + + # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler._sigma_to_t + def _sigma_to_t(self, sigma): + return sigma * self.config.num_train_timesteps + + def _sigma_to_alpha_sigma_t(self, sigma): + return 1 - sigma, sigma + + # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.set_timesteps + def time_shift(self, mu: float, sigma: float, t: torch.Tensor): + return math.exp(mu) / (math.exp(mu) + (1 / t - 1)**sigma) + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.convert_model_output + def convert_model_output( + self, + model_output: torch.Tensor, + *args, + sample: torch.Tensor = None, + **kwargs, + ) -> torch.Tensor: + """ + Convert the model output to the corresponding type the DPMSolver/DPMSolver++ algorithm needs. DPM-Solver is + designed to discretize an integral of the noise prediction model, and DPM-Solver++ is designed to discretize an + integral of the data prediction model. + + The algorithm and model type are decoupled. You can use either DPMSolver or DPMSolver++ for both noise + prediction and data prediction models. + + Args: + model_output (`torch.Tensor`): + The direct output from the learned diffusion model. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + Returns: + `torch.Tensor`: + The converted model output. + """ + timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None) + if sample is None: + if len(args) > 1: + sample = args[1] + else: + raise ValueError( + "missing `sample` as a required keyward argument") + if timestep is not None: + deprecate( + "timesteps", + "1.0.0", + "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", + ) + + # DPM-Solver++ needs to solve an integral of the data prediction model. + if self.config.algorithm_type in ["dpmsolver++", "sde-dpmsolver++"]: + if self.config.prediction_type == "flow_prediction": + sigma_t = self.sigmas[self.step_index] + x0_pred = sample - sigma_t * model_output + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`," + " `v_prediction`, or `flow_prediction` for the FlowDPMSolverMultistepScheduler." + ) + + if self.config.thresholding: + x0_pred = self._threshold_sample(x0_pred) + + return x0_pred + + # DPM-Solver needs to solve an integral of the noise prediction model. + elif self.config.algorithm_type in ["dpmsolver", "sde-dpmsolver"]: + if self.config.prediction_type == "flow_prediction": + sigma_t = self.sigmas[self.step_index] + epsilon = sample - (1 - sigma_t) * model_output + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`," + " `v_prediction` or `flow_prediction` for the FlowDPMSolverMultistepScheduler." + ) + + if self.config.thresholding: + sigma_t = self.sigmas[self.step_index] + x0_pred = sample - sigma_t * model_output + x0_pred = self._threshold_sample(x0_pred) + epsilon = model_output + x0_pred + + return epsilon + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.dpm_solver_first_order_update + def dpm_solver_first_order_update( + self, + model_output: torch.Tensor, + *args, + sample: torch.Tensor = None, + noise: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + """ + One step for the first-order DPMSolver (equivalent to DDIM). + Args: + model_output (`torch.Tensor`): + The direct output from the learned diffusion model. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + Returns: + `torch.Tensor`: + The sample tensor at the previous timestep. + """ + timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None) + prev_timestep = args[1] if len(args) > 1 else kwargs.pop( + "prev_timestep", None) + if sample is None: + if len(args) > 2: + sample = args[2] + else: + raise ValueError( + " missing `sample` as a required keyward argument") + if timestep is not None: + deprecate( + "timesteps", + "1.0.0", + "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", + ) + + if prev_timestep is not None: + deprecate( + "prev_timestep", + "1.0.0", + "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", + ) + + sigma_t, sigma_s = self.sigmas[self.step_index + 1], self.sigmas[ + self.step_index] # pyright: ignore + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s, sigma_s = self._sigma_to_alpha_sigma_t(sigma_s) + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s = torch.log(alpha_s) - torch.log(sigma_s) + + h = lambda_t - lambda_s + if self.config.algorithm_type == "dpmsolver++": + x_t = (sigma_t / + sigma_s) * sample - (alpha_t * + (torch.exp(-h) - 1.0)) * model_output + elif self.config.algorithm_type == "dpmsolver": + x_t = (alpha_t / + alpha_s) * sample - (sigma_t * + (torch.exp(h) - 1.0)) * model_output + elif self.config.algorithm_type == "sde-dpmsolver++": + assert noise is not None + x_t = ((sigma_t / sigma_s * torch.exp(-h)) * sample + + (alpha_t * (1 - torch.exp(-2.0 * h))) * model_output + + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise) + elif self.config.algorithm_type == "sde-dpmsolver": + assert noise is not None + x_t = ((alpha_t / alpha_s) * sample - 2.0 * + (sigma_t * (torch.exp(h) - 1.0)) * model_output + + sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise) + return x_t # pyright: ignore + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.multistep_dpm_solver_second_order_update + def multistep_dpm_solver_second_order_update( + self, + model_output_list: List[torch.Tensor], + *args, + sample: torch.Tensor = None, + noise: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + """ + One step for the second-order multistep DPMSolver. + Args: + model_output_list (`List[torch.Tensor]`): + The direct outputs from learned diffusion model at current and latter timesteps. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + Returns: + `torch.Tensor`: + The sample tensor at the previous timestep. + """ + timestep_list = args[0] if len(args) > 0 else kwargs.pop( + "timestep_list", None) + prev_timestep = args[1] if len(args) > 1 else kwargs.pop( + "prev_timestep", None) + if sample is None: + if len(args) > 2: + sample = args[2] + else: + raise ValueError( + " missing `sample` as a required keyward argument") + if timestep_list is not None: + deprecate( + "timestep_list", + "1.0.0", + "Passing `timestep_list` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", + ) + + if prev_timestep is not None: + deprecate( + "prev_timestep", + "1.0.0", + "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", + ) + + sigma_t, sigma_s0, sigma_s1 = ( + self.sigmas[self.step_index + 1], # pyright: ignore + self.sigmas[self.step_index], + self.sigmas[self.step_index - 1], # pyright: ignore + ) + + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1) + + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1) + + m0, m1 = model_output_list[-1], model_output_list[-2] + + h, h_0 = lambda_t - lambda_s0, lambda_s0 - lambda_s1 + r0 = h_0 / h + D0, D1 = m0, (1.0 / r0) * (m0 - m1) + if self.config.algorithm_type == "dpmsolver++": + # See https://arxiv.org/abs/2211.01095 for detailed derivations + if self.config.solver_type == "midpoint": + x_t = ((sigma_t / sigma_s0) * sample - + (alpha_t * (torch.exp(-h) - 1.0)) * D0 - 0.5 * + (alpha_t * (torch.exp(-h) - 1.0)) * D1) + elif self.config.solver_type == "heun": + x_t = ((sigma_t / sigma_s0) * sample - + (alpha_t * (torch.exp(-h) - 1.0)) * D0 + + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1) + elif self.config.algorithm_type == "dpmsolver": + # See https://arxiv.org/abs/2206.00927 for detailed derivations + if self.config.solver_type == "midpoint": + x_t = ((alpha_t / alpha_s0) * sample - + (sigma_t * (torch.exp(h) - 1.0)) * D0 - 0.5 * + (sigma_t * (torch.exp(h) - 1.0)) * D1) + elif self.config.solver_type == "heun": + x_t = ((alpha_t / alpha_s0) * sample - + (sigma_t * (torch.exp(h) - 1.0)) * D0 - + (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1) + elif self.config.algorithm_type == "sde-dpmsolver++": + assert noise is not None + if self.config.solver_type == "midpoint": + x_t = ((sigma_t / sigma_s0 * torch.exp(-h)) * sample + + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0 + 0.5 * + (alpha_t * (1 - torch.exp(-2.0 * h))) * D1 + + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise) + elif self.config.solver_type == "heun": + x_t = ((sigma_t / sigma_s0 * torch.exp(-h)) * sample + + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0 + + (alpha_t * ((1.0 - torch.exp(-2.0 * h)) / + (-2.0 * h) + 1.0)) * D1 + + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise) + elif self.config.algorithm_type == "sde-dpmsolver": + assert noise is not None + if self.config.solver_type == "midpoint": + x_t = ((alpha_t / alpha_s0) * sample - 2.0 * + (sigma_t * (torch.exp(h) - 1.0)) * D0 - + (sigma_t * (torch.exp(h) - 1.0)) * D1 + + sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise) + elif self.config.solver_type == "heun": + x_t = ((alpha_t / alpha_s0) * sample - 2.0 * + (sigma_t * (torch.exp(h) - 1.0)) * D0 - 2.0 * + (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1 + + sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise) + return x_t # pyright: ignore + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.multistep_dpm_solver_third_order_update + def multistep_dpm_solver_third_order_update( + self, + model_output_list: List[torch.Tensor], + *args, + sample: torch.Tensor = None, + **kwargs, + ) -> torch.Tensor: + """ + One step for the third-order multistep DPMSolver. + Args: + model_output_list (`List[torch.Tensor]`): + The direct outputs from learned diffusion model at current and latter timesteps. + sample (`torch.Tensor`): + A current instance of a sample created by diffusion process. + Returns: + `torch.Tensor`: + The sample tensor at the previous timestep. + """ + + timestep_list = args[0] if len(args) > 0 else kwargs.pop( + "timestep_list", None) + prev_timestep = args[1] if len(args) > 1 else kwargs.pop( + "prev_timestep", None) + if sample is None: + if len(args) > 2: + sample = args[2] + else: + raise ValueError( + " missing`sample` as a required keyward argument") + if timestep_list is not None: + deprecate( + "timestep_list", + "1.0.0", + "Passing `timestep_list` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", + ) + + if prev_timestep is not None: + deprecate( + "prev_timestep", + "1.0.0", + "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", + ) + + sigma_t, sigma_s0, sigma_s1, sigma_s2 = ( + self.sigmas[self.step_index + 1], # pyright: ignore + self.sigmas[self.step_index], + self.sigmas[self.step_index - 1], # pyright: ignore + self.sigmas[self.step_index - 2], # pyright: ignore + ) + + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1) + alpha_s2, sigma_s2 = self._sigma_to_alpha_sigma_t(sigma_s2) + + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1) + lambda_s2 = torch.log(alpha_s2) - torch.log(sigma_s2) + + m0, m1, m2 = model_output_list[-1], model_output_list[ + -2], model_output_list[-3] + + h, h_0, h_1 = lambda_t - lambda_s0, lambda_s0 - lambda_s1, lambda_s1 - lambda_s2 + r0, r1 = h_0 / h, h_1 / h + D0 = m0 + D1_0, D1_1 = (1.0 / r0) * (m0 - m1), (1.0 / r1) * (m1 - m2) + D1 = D1_0 + (r0 / (r0 + r1)) * (D1_0 - D1_1) + D2 = (1.0 / (r0 + r1)) * (D1_0 - D1_1) + if self.config.algorithm_type == "dpmsolver++": + # See https://arxiv.org/abs/2206.00927 for detailed derivations + x_t = ((sigma_t / sigma_s0) * sample - + (alpha_t * (torch.exp(-h) - 1.0)) * D0 + + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1 - + (alpha_t * ((torch.exp(-h) - 1.0 + h) / h**2 - 0.5)) * D2) + elif self.config.algorithm_type == "dpmsolver": + # See https://arxiv.org/abs/2206.00927 for detailed derivations + x_t = ((alpha_t / alpha_s0) * sample - (sigma_t * + (torch.exp(h) - 1.0)) * D0 - + (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1 - + (sigma_t * ((torch.exp(h) - 1.0 - h) / h**2 - 0.5)) * D2) + return x_t # pyright: ignore + + def index_for_timestep(self, timestep, schedule_timesteps=None): + if schedule_timesteps is None: + schedule_timesteps = self.timesteps + + indices = (schedule_timesteps == timestep).nonzero() + + # The sigma index that is taken for the **very** first `step` + # is always the second index (or the last index if there is only 1) + # This way we can ensure we don't accidentally skip a sigma in + # case we start in the middle of the denoising schedule (e.g. for image-to-image) + pos = 1 if len(indices) > 1 else 0 + + return indices[pos].item() + + def _init_step_index(self, timestep): + """ + Initialize the step_index counter for the scheduler. + """ + + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + # Modified from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.step + def step( + self, + model_output: torch.Tensor, + timestep: Union[int, torch.Tensor], + sample: torch.Tensor, + generator=None, + variance_noise: Optional[torch.Tensor] = None, + return_dict: bool = True, + ) -> Union[SchedulerOutput, Tuple]: + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with + the multistep DPMSolver. + Args: + model_output (`torch.Tensor`): + The direct output from learned diffusion model. + timestep (`int`): + The current discrete timestep in the diffusion chain. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + generator (`torch.Generator`, *optional*): + A random number generator. + variance_noise (`torch.Tensor`): + Alternative to generating noise with `generator` by directly providing the noise for the variance + itself. Useful for methods such as [`LEdits++`]. + return_dict (`bool`): + Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`. + Returns: + [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`: + If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a + tuple is returned where the first element is the sample tensor. + """ + if self.num_inference_steps is None: + raise ValueError( + "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" + ) + + if self.step_index is None: + self._init_step_index(timestep) + + # Improve numerical stability for small number of steps + lower_order_final = (self.step_index == len(self.timesteps) - 1) and ( + self.config.euler_at_final or + (self.config.lower_order_final and len(self.timesteps) < 15) or + self.config.final_sigmas_type == "zero") + lower_order_second = ((self.step_index == len(self.timesteps) - 2) and + self.config.lower_order_final and + len(self.timesteps) < 15) + + model_output = self.convert_model_output(model_output, sample=sample) + for i in range(self.config.solver_order - 1): + self.model_outputs[i] = self.model_outputs[i + 1] + self.model_outputs[-1] = model_output + + # Upcast to avoid precision issues when computing prev_sample + sample = sample.to(torch.float32) + if self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++" + ] and variance_noise is None: + noise = randn_tensor( + model_output.shape, + generator=generator, + device=model_output.device, + dtype=torch.float32) + elif self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"]: + noise = variance_noise.to( + device=model_output.device, + dtype=torch.float32) # pyright: ignore + else: + noise = None + + if self.config.solver_order == 1 or self.lower_order_nums < 1 or lower_order_final: + prev_sample = self.dpm_solver_first_order_update( + model_output, sample=sample, noise=noise) + elif self.config.solver_order == 2 or self.lower_order_nums < 2 or lower_order_second: + prev_sample = self.multistep_dpm_solver_second_order_update( + self.model_outputs, sample=sample, noise=noise) + else: + prev_sample = self.multistep_dpm_solver_third_order_update( + self.model_outputs, sample=sample) + + if self.lower_order_nums < self.config.solver_order: + self.lower_order_nums += 1 + + # Cast sample back to expected dtype + prev_sample = prev_sample.to(model_output.dtype) + + # upon completion increase step index by one + self._step_index += 1 # pyright: ignore + + if not return_dict: + return (prev_sample,) + + return SchedulerOutput(prev_sample=prev_sample) + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.scale_model_input + def scale_model_input(self, sample: torch.Tensor, *args, + **kwargs) -> torch.Tensor: + """ + Ensures interchangeability with schedulers that need to scale the denoising model input depending on the + current timestep. + Args: + sample (`torch.Tensor`): + The input sample. + Returns: + `torch.Tensor`: + A scaled input sample. + """ + return sample + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.scale_model_input + def add_noise( + self, + original_samples: torch.Tensor, + noise: torch.Tensor, + timesteps: torch.IntTensor, + ) -> torch.Tensor: + # Make sure sigmas and timesteps have the same device and dtype as original_samples + sigmas = self.sigmas.to( + device=original_samples.device, dtype=original_samples.dtype) + if original_samples.device.type == "mps" and torch.is_floating_point( + timesteps): + # mps does not support float64 + schedule_timesteps = self.timesteps.to( + original_samples.device, dtype=torch.float32) + timesteps = timesteps.to( + original_samples.device, dtype=torch.float32) + else: + schedule_timesteps = self.timesteps.to(original_samples.device) + timesteps = timesteps.to(original_samples.device) + + # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index + if self.begin_index is None: + step_indices = [ + self.index_for_timestep(t, schedule_timesteps) + for t in timesteps + ] + elif self.step_index is not None: + # add_noise is called after first denoising step (for inpainting) + step_indices = [self.step_index] * timesteps.shape[0] + else: + # add noise is called before first denoising step to create initial latent(img2img) + step_indices = [self.begin_index] * timesteps.shape[0] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < len(original_samples.shape): + sigma = sigma.unsqueeze(-1) + + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) + noisy_samples = alpha_t * original_samples + sigma_t * noise + return noisy_samples + + def __len__(self): + return self.config.num_train_timesteps diff --git a/wan_5b/utils/fm_solvers_unipc.py b/wan_5b/utils/fm_solvers_unipc.py new file mode 100644 index 0000000..fb502f2 --- /dev/null +++ b/wan_5b/utils/fm_solvers_unipc.py @@ -0,0 +1,802 @@ +# Copied from https://github.com/huggingface/diffusers/blob/v0.31.0/src/diffusers/schedulers/scheduling_unipc_multistep.py +# Convert unipc for flow matching +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. + +import math +from typing import List, Optional, Tuple, Union + +import numpy as np +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import ( + KarrasDiffusionSchedulers, + SchedulerMixin, + SchedulerOutput, +) +from diffusers.utils import deprecate, is_scipy_available + +if is_scipy_available(): + import scipy.stats + + +class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin): + """ + `UniPCMultistepScheduler` is a training-free framework designed for the fast sampling of diffusion models. + + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps to train the model. + solver_order (`int`, default `2`): + The UniPC order which can be any positive integer. The effective order of accuracy is `solver_order + 1` + due to the UniC. It is recommended to use `solver_order=2` for guided sampling, and `solver_order=3` for + unconditional sampling. + prediction_type (`str`, defaults to "flow_prediction"): + Prediction type of the scheduler function; must be `flow_prediction` for this scheduler, which predicts + the flow of the diffusion process. + thresholding (`bool`, defaults to `False`): + Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such + as Stable Diffusion. + dynamic_thresholding_ratio (`float`, defaults to 0.995): + The ratio for the dynamic thresholding method. Valid only when `thresholding=True`. + sample_max_value (`float`, defaults to 1.0): + The threshold value for dynamic thresholding. Valid only when `thresholding=True` and `predict_x0=True`. + predict_x0 (`bool`, defaults to `True`): + Whether to use the updating algorithm on the predicted x0. + solver_type (`str`, default `bh2`): + Solver type for UniPC. It is recommended to use `bh1` for unconditional sampling when steps < 10, and `bh2` + otherwise. + lower_order_final (`bool`, default `True`): + Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can + stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10. + disable_corrector (`list`, default `[]`): + Decides which step to disable the corrector to mitigate the misalignment between `epsilon_theta(x_t, c)` + and `epsilon_theta(x_t^c, c)` which can influence convergence for a large guidance scale. Corrector is + usually disabled during the first few steps. + solver_p (`SchedulerMixin`, default `None`): + Any other scheduler that if specified, the algorithm becomes `solver_p + UniC`. + use_karras_sigmas (`bool`, *optional*, defaults to `False`): + Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`, + the sigmas are determined according to a sequence of noise levels {σi}. + use_exponential_sigmas (`bool`, *optional*, defaults to `False`): + Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process. + timestep_spacing (`str`, defaults to `"linspace"`): + The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. + steps_offset (`int`, defaults to 0): + An offset added to the inference steps, as required by some model families. + final_sigmas_type (`str`, defaults to `"zero"`): + The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final + sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0. + """ + + _compatibles = [e.name for e in KarrasDiffusionSchedulers] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + solver_order: int = 2, + prediction_type: str = "flow_prediction", + shift: Optional[float] = 1.0, + use_dynamic_shifting=False, + thresholding: bool = False, + dynamic_thresholding_ratio: float = 0.995, + sample_max_value: float = 1.0, + predict_x0: bool = True, + solver_type: str = "bh2", + lower_order_final: bool = True, + disable_corrector: List[int] = [], + solver_p: SchedulerMixin = None, + timestep_spacing: str = "linspace", + steps_offset: int = 0, + final_sigmas_type: Optional[str] = "zero", # "zero", "sigma_min" + ): + + if solver_type not in ["bh1", "bh2"]: + if solver_type in ["midpoint", "heun", "logrho"]: + self.register_to_config(solver_type="bh2") + else: + raise NotImplementedError( + f"{solver_type} is not implemented for {self.__class__}") + + self.predict_x0 = predict_x0 + # setable values + self.num_inference_steps = None + alphas = np.linspace(1, 1 / num_train_timesteps, + num_train_timesteps)[::-1].copy() + sigmas = 1.0 - alphas + sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32) + + if not use_dynamic_shifting: + # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution + sigmas = shift * sigmas / (1 + + (shift - 1) * sigmas) # pyright: ignore + + self.sigmas = sigmas + self.timesteps = sigmas * num_train_timesteps + + self.model_outputs = [None] * solver_order + self.timestep_list = [None] * solver_order + self.lower_order_nums = 0 + self.disable_corrector = disable_corrector + self.solver_p = solver_p + self.last_sample = None + self._step_index = None + self._begin_index = None + + self.sigmas = self.sigmas.to( + "cpu") # to avoid too much CPU/GPU communication + self.sigma_min = self.sigmas[-1].item() + self.sigma_max = self.sigmas[0].item() + + @property + def step_index(self): + """ + The index counter for current timestep. It will increase 1 after each scheduler step. + """ + return self._step_index + + @property + def begin_index(self): + """ + The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + """ + return self._begin_index + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index + def set_begin_index(self, begin_index: int = 0): + """ + Sets the begin index for the scheduler. This function should be run from pipeline before the inference. + + Args: + begin_index (`int`): + The begin index for the scheduler. + """ + self._begin_index = begin_index + + # Modified from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.set_timesteps + def set_timesteps( + self, + num_inference_steps: Union[int, None] = None, + device: Union[str, torch.device] = None, + sigmas: Optional[List[float]] = None, + mu: Optional[Union[float, None]] = None, + shift: Optional[Union[float, None]] = None, + ): + """ + Sets the discrete timesteps used for the diffusion chain (to be run before inference). + Args: + num_inference_steps (`int`): + Total number of the spacing of the time steps. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + """ + + if self.config.use_dynamic_shifting and mu is None: + raise ValueError( + " you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`" + ) + + if sigmas is None: + sigmas = np.linspace(self.sigma_max, self.sigma_min, + num_inference_steps + + 1).copy()[:-1] # pyright: ignore + + if self.config.use_dynamic_shifting: + sigmas = self.time_shift(mu, 1.0, sigmas) # pyright: ignore + else: + if shift is None: + shift = self.config.shift + sigmas = shift * sigmas / (1 + + (shift - 1) * sigmas) # pyright: ignore + + if self.config.final_sigmas_type == "sigma_min": + sigma_last = ((1 - self.alphas_cumprod[0]) / + self.alphas_cumprod[0])**0.5 + elif self.config.final_sigmas_type == "zero": + sigma_last = 0 + else: + raise ValueError( + f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" + ) + + timesteps = sigmas * self.config.num_train_timesteps + sigmas = np.concatenate([sigmas, [sigma_last] + ]).astype(np.float32) # pyright: ignore + + self.sigmas = torch.from_numpy(sigmas) + self.timesteps = torch.from_numpy(timesteps).to( + device=device, dtype=torch.int64) + + self.num_inference_steps = len(timesteps) + + self.model_outputs = [ + None, + ] * self.config.solver_order + self.lower_order_nums = 0 + self.last_sample = None + if self.solver_p: + self.solver_p.set_timesteps(self.num_inference_steps, device=device) + + # add an index counter for schedulers that allow duplicated timesteps + self._step_index = None + self._begin_index = None + self.sigmas = self.sigmas.to( + "cpu") # to avoid too much CPU/GPU communication + + # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample + def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: + """ + "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the + prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by + s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing + pixels from saturation at each step. We find that dynamic thresholding results in significantly better + photorealism as well as better image-text alignment, especially when using very large guidance weights." + + https://arxiv.org/abs/2205.11487 + """ + dtype = sample.dtype + batch_size, channels, *remaining_dims = sample.shape + + if dtype not in (torch.float32, torch.float64): + sample = sample.float( + ) # upcast for quantile calculation, and clamp not implemented for cpu half + + # Flatten sample for doing quantile calculation along each image + sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) + + abs_sample = sample.abs() # "a certain percentile absolute pixel value" + + s = torch.quantile( + abs_sample, self.config.dynamic_thresholding_ratio, dim=1) + s = torch.clamp( + s, min=1, max=self.config.sample_max_value + ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] + s = s.unsqueeze( + 1) # (batch_size, 1) because clamp will broadcast along dim=0 + sample = torch.clamp( + sample, -s, s + ) / s # "we threshold xt0 to the range [-s, s] and then divide by s" + + sample = sample.reshape(batch_size, channels, *remaining_dims) + sample = sample.to(dtype) + + return sample + + # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler._sigma_to_t + def _sigma_to_t(self, sigma): + return sigma * self.config.num_train_timesteps + + def _sigma_to_alpha_sigma_t(self, sigma): + return 1 - sigma, sigma + + # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.set_timesteps + def time_shift(self, mu: float, sigma: float, t: torch.Tensor): + return math.exp(mu) / (math.exp(mu) + (1 / t - 1)**sigma) + + def convert_model_output( + self, + model_output: torch.Tensor, + *args, + sample: torch.Tensor = None, + **kwargs, + ) -> torch.Tensor: + r""" + Convert the model output to the corresponding type the UniPC algorithm needs. + + Args: + model_output (`torch.Tensor`): + The direct output from the learned diffusion model. + timestep (`int`): + The current discrete timestep in the diffusion chain. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + + Returns: + `torch.Tensor`: + The converted model output. + """ + timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None) + if sample is None: + if len(args) > 1: + sample = args[1] + else: + raise ValueError( + "missing `sample` as a required keyward argument") + if timestep is not None: + deprecate( + "timesteps", + "1.0.0", + "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", + ) + + sigma = self.sigmas[self.step_index] + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) + + if self.predict_x0: + if self.config.prediction_type == "flow_prediction": + sigma_t = self.sigmas[self.step_index] + x0_pred = sample - sigma_t * model_output + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`," + " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler." + ) + + if self.config.thresholding: + x0_pred = self._threshold_sample(x0_pred) + + return x0_pred + else: + if self.config.prediction_type == "flow_prediction": + sigma_t = self.sigmas[self.step_index] + epsilon = sample - (1 - sigma_t) * model_output + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`," + " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler." + ) + + if self.config.thresholding: + sigma_t = self.sigmas[self.step_index] + x0_pred = sample - sigma_t * model_output + x0_pred = self._threshold_sample(x0_pred) + epsilon = model_output + x0_pred + + return epsilon + + def multistep_uni_p_bh_update( + self, + model_output: torch.Tensor, + *args, + sample: torch.Tensor = None, + order: int = None, # pyright: ignore + **kwargs, + ) -> torch.Tensor: + """ + One step for the UniP (B(h) version). Alternatively, `self.solver_p` is used if is specified. + + Args: + model_output (`torch.Tensor`): + The direct output from the learned diffusion model at the current timestep. + prev_timestep (`int`): + The previous discrete timestep in the diffusion chain. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + order (`int`): + The order of UniP at this timestep (corresponds to the *p* in UniPC-p). + + Returns: + `torch.Tensor`: + The sample tensor at the previous timestep. + """ + prev_timestep = args[0] if len(args) > 0 else kwargs.pop( + "prev_timestep", None) + if sample is None: + if len(args) > 1: + sample = args[1] + else: + raise ValueError( + " missing `sample` as a required keyward argument") + if order is None: + if len(args) > 2: + order = args[2] + else: + raise ValueError( + " missing `order` as a required keyward argument") + if prev_timestep is not None: + deprecate( + "prev_timestep", + "1.0.0", + "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", + ) + model_output_list = self.model_outputs + + s0 = self.timestep_list[-1] + m0 = model_output_list[-1] + x = sample + + if self.solver_p: + x_t = self.solver_p.step(model_output, s0, x).prev_sample + return x_t + + sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[ + self.step_index] # pyright: ignore + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + + h = lambda_t - lambda_s0 + device = sample.device + + rks = [] + D1s = [] + for i in range(1, order): + si = self.step_index - i # pyright: ignore + mi = model_output_list[-(i + 1)] + alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) + lambda_si = torch.log(alpha_si) - torch.log(sigma_si) + rk = (lambda_si - lambda_s0) / h + rks.append(rk) + D1s.append((mi - m0) / rk) # pyright: ignore + + rks.append(1.0) + rks = torch.tensor(rks, device=device) + + R = [] + b = [] + + hh = -h if self.predict_x0 else h + h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1 + h_phi_k = h_phi_1 / hh - 1 + + factorial_i = 1 + + if self.config.solver_type == "bh1": + B_h = hh + elif self.config.solver_type == "bh2": + B_h = torch.expm1(hh) + else: + raise NotImplementedError() + + for i in range(1, order + 1): + R.append(torch.pow(rks, i - 1)) + b.append(h_phi_k * factorial_i / B_h) + factorial_i *= i + 1 + h_phi_k = h_phi_k / hh - 1 / factorial_i + + R = torch.stack(R) + b = torch.tensor(b, device=device) + + if len(D1s) > 0: + D1s = torch.stack(D1s, dim=1) # (B, K) + # for order 2, we use a simplified version + if order == 2: + rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device) + else: + rhos_p = torch.linalg.solve(R[:-1, :-1], + b[:-1]).to(device).to(x.dtype) + else: + D1s = None + + if self.predict_x0: + x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 + if D1s is not None: + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, + D1s) # pyright: ignore + else: + pred_res = 0 + x_t = x_t_ - alpha_t * B_h * pred_res + else: + x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 + if D1s is not None: + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, + D1s) # pyright: ignore + else: + pred_res = 0 + x_t = x_t_ - sigma_t * B_h * pred_res + + x_t = x_t.to(x.dtype) + return x_t + + def multistep_uni_c_bh_update( + self, + this_model_output: torch.Tensor, + *args, + last_sample: torch.Tensor = None, + this_sample: torch.Tensor = None, + order: int = None, # pyright: ignore + **kwargs, + ) -> torch.Tensor: + """ + One step for the UniC (B(h) version). + + Args: + this_model_output (`torch.Tensor`): + The model outputs at `x_t`. + this_timestep (`int`): + The current timestep `t`. + last_sample (`torch.Tensor`): + The generated sample before the last predictor `x_{t-1}`. + this_sample (`torch.Tensor`): + The generated sample after the last predictor `x_{t}`. + order (`int`): + The `p` of UniC-p at this step. The effective order of accuracy should be `order + 1`. + + Returns: + `torch.Tensor`: + The corrected sample tensor at the current timestep. + """ + this_timestep = args[0] if len(args) > 0 else kwargs.pop( + "this_timestep", None) + if last_sample is None: + if len(args) > 1: + last_sample = args[1] + else: + raise ValueError( + " missing`last_sample` as a required keyward argument") + if this_sample is None: + if len(args) > 2: + this_sample = args[2] + else: + raise ValueError( + " missing`this_sample` as a required keyward argument") + if order is None: + if len(args) > 3: + order = args[3] + else: + raise ValueError( + " missing`order` as a required keyward argument") + if this_timestep is not None: + deprecate( + "this_timestep", + "1.0.0", + "Passing `this_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`", + ) + + model_output_list = self.model_outputs + + m0 = model_output_list[-1] + x = last_sample + x_t = this_sample + model_t = this_model_output + + sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[ + self.step_index - 1] # pyright: ignore + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + + h = lambda_t - lambda_s0 + device = this_sample.device + + rks = [] + D1s = [] + for i in range(1, order): + si = self.step_index - (i + 1) # pyright: ignore + mi = model_output_list[-(i + 1)] + alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) + lambda_si = torch.log(alpha_si) - torch.log(sigma_si) + rk = (lambda_si - lambda_s0) / h + rks.append(rk) + D1s.append((mi - m0) / rk) # pyright: ignore + + rks.append(1.0) + rks = torch.tensor(rks, device=device) + + R = [] + b = [] + + hh = -h if self.predict_x0 else h + h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1 + h_phi_k = h_phi_1 / hh - 1 + + factorial_i = 1 + + if self.config.solver_type == "bh1": + B_h = hh + elif self.config.solver_type == "bh2": + B_h = torch.expm1(hh) + else: + raise NotImplementedError() + + for i in range(1, order + 1): + R.append(torch.pow(rks, i - 1)) + b.append(h_phi_k * factorial_i / B_h) + factorial_i *= i + 1 + h_phi_k = h_phi_k / hh - 1 / factorial_i + + R = torch.stack(R) + b = torch.tensor(b, device=device) + + if len(D1s) > 0: + D1s = torch.stack(D1s, dim=1) + else: + D1s = None + + # for order 1, we use a simplified version + if order == 1: + rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device) + else: + rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype) + + if self.predict_x0: + x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 + if D1s is not None: + corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) + else: + corr_res = 0 + D1_t = model_t - m0 + x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t) + else: + x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 + if D1s is not None: + corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) + else: + corr_res = 0 + D1_t = model_t - m0 + x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t) + x_t = x_t.to(x.dtype) + return x_t + + def index_for_timestep(self, timestep, schedule_timesteps=None): + if schedule_timesteps is None: + schedule_timesteps = self.timesteps + + indices = (schedule_timesteps == timestep).nonzero() + + # The sigma index that is taken for the **very** first `step` + # is always the second index (or the last index if there is only 1) + # This way we can ensure we don't accidentally skip a sigma in + # case we start in the middle of the denoising schedule (e.g. for image-to-image) + pos = 1 if len(indices) > 1 else 0 + + return indices[pos].item() + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index + def _init_step_index(self, timestep): + """ + Initialize the step_index counter for the scheduler. + """ + + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + def step(self, + model_output: torch.Tensor, + timestep: Union[int, torch.Tensor], + sample: torch.Tensor, + return_dict: bool = True, + generator=None) -> Union[SchedulerOutput, Tuple]: + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with + the multistep UniPC. + + Args: + model_output (`torch.Tensor`): + The direct output from learned diffusion model. + timestep (`int`): + The current discrete timestep in the diffusion chain. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + return_dict (`bool`): + Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`. + + Returns: + [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`: + If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a + tuple is returned where the first element is the sample tensor. + + """ + if self.num_inference_steps is None: + raise ValueError( + "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" + ) + + if self.step_index is None: + self._init_step_index(timestep) + + use_corrector = ( + self.step_index > 0 and + self.step_index - 1 not in self.disable_corrector and + self.last_sample is not None # pyright: ignore + ) + + model_output_convert = self.convert_model_output( + model_output, sample=sample) + if use_corrector: + sample = self.multistep_uni_c_bh_update( + this_model_output=model_output_convert, + last_sample=self.last_sample, + this_sample=sample, + order=self.this_order, + ) + + for i in range(self.config.solver_order - 1): + self.model_outputs[i] = self.model_outputs[i + 1] + self.timestep_list[i] = self.timestep_list[i + 1] + + self.model_outputs[-1] = model_output_convert + self.timestep_list[-1] = timestep # pyright: ignore + + if self.config.lower_order_final: + this_order = min(self.config.solver_order, + len(self.timesteps) - + self.step_index) # pyright: ignore + else: + this_order = self.config.solver_order + + self.this_order = min(this_order, + self.lower_order_nums + 1) # warmup for multistep + assert self.this_order > 0 + + self.last_sample = sample + prev_sample = self.multistep_uni_p_bh_update( + model_output=model_output, # pass the original non-converted model output, in case solver-p is used + sample=sample, + order=self.this_order, + ) + + if self.lower_order_nums < self.config.solver_order: + self.lower_order_nums += 1 + + # upon completion increase step index by one + self._step_index += 1 # pyright: ignore + + if not return_dict: + return (prev_sample,) + + return SchedulerOutput(prev_sample=prev_sample) + + def scale_model_input(self, sample: torch.Tensor, *args, + **kwargs) -> torch.Tensor: + """ + Ensures interchangeability with schedulers that need to scale the denoising model input depending on the + current timestep. + + Args: + sample (`torch.Tensor`): + The input sample. + + Returns: + `torch.Tensor`: + A scaled input sample. + """ + return sample + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.add_noise + def add_noise( + self, + original_samples: torch.Tensor, + noise: torch.Tensor, + timesteps: torch.IntTensor, + ) -> torch.Tensor: + # Make sure sigmas and timesteps have the same device and dtype as original_samples + sigmas = self.sigmas.to( + device=original_samples.device, dtype=original_samples.dtype) + if original_samples.device.type == "mps" and torch.is_floating_point( + timesteps): + # mps does not support float64 + schedule_timesteps = self.timesteps.to( + original_samples.device, dtype=torch.float32) + timesteps = timesteps.to( + original_samples.device, dtype=torch.float32) + else: + schedule_timesteps = self.timesteps.to(original_samples.device) + timesteps = timesteps.to(original_samples.device) + + # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index + if self.begin_index is None: + step_indices = [ + self.index_for_timestep(t, schedule_timesteps) + for t in timesteps + ] + elif self.step_index is not None: + # add_noise is called after first denoising step (for inpainting) + step_indices = [self.step_index] * timesteps.shape[0] + else: + # add noise is called before first denoising step to create initial latent(img2img) + step_indices = [self.begin_index] * timesteps.shape[0] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < len(original_samples.shape): + sigma = sigma.unsqueeze(-1) + + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) + noisy_samples = alpha_t * original_samples + sigma_t * noise + return noisy_samples + + def __len__(self): + return self.config.num_train_timesteps diff --git a/wan_5b/utils/prompt_extend.py b/wan_5b/utils/prompt_extend.py new file mode 100644 index 0000000..9d40d9c --- /dev/null +++ b/wan_5b/utils/prompt_extend.py @@ -0,0 +1,542 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import json +import logging +import math +import os +import random +import sys +import tempfile +from dataclasses import dataclass +from http import HTTPStatus +from typing import Optional, Union + +import dashscope +import torch +from PIL import Image + +try: + from flash_attn import flash_attn_varlen_func + FLASH_VER = 2 +except ModuleNotFoundError: + flash_attn_varlen_func = None # in compatible with CPU machines + FLASH_VER = None + +from .system_prompt import * + +DEFAULT_SYS_PROMPTS = { + "t2v-A14B": { + "zh": T2V_A14B_ZH_SYS_PROMPT, + "en": T2V_A14B_EN_SYS_PROMPT, + }, + "i2v-A14B": { + "zh": I2V_A14B_ZH_SYS_PROMPT, + "en": I2V_A14B_EN_SYS_PROMPT, + "empty": { + "zh": I2V_A14B_EMPTY_ZH_SYS_PROMPT, + "en": I2V_A14B_EMPTY_EN_SYS_PROMPT, + } + }, + "ti2v-5B": { + "t2v": { + "zh": T2V_A14B_ZH_SYS_PROMPT, + "en": T2V_A14B_EN_SYS_PROMPT, + }, + "i2v": { + "zh": I2V_A14B_ZH_SYS_PROMPT, + "en": I2V_A14B_EN_SYS_PROMPT, + } + }, +} + + +@dataclass +class PromptOutput(object): + status: bool + prompt: str + seed: int + system_prompt: str + message: str + + def add_custom_field(self, key: str, value) -> None: + self.__setattr__(key, value) + + +class PromptExpander: + + def __init__(self, model_name, task, is_vl=False, device=0, **kwargs): + self.model_name = model_name + self.task = task + self.is_vl = is_vl + self.device = device + + def extend_with_img(self, + prompt, + system_prompt, + image=None, + seed=-1, + *args, + **kwargs): + pass + + def extend(self, prompt, system_prompt, seed=-1, *args, **kwargs): + pass + + def decide_system_prompt(self, tar_lang="zh", prompt=None): + assert self.task is not None + if "ti2v" in self.task: + if self.is_vl: + return DEFAULT_SYS_PROMPTS[self.task]["i2v"][tar_lang] + else: + return DEFAULT_SYS_PROMPTS[self.task]["t2v"][tar_lang] + if "i2v" in self.task and len(prompt) == 0: + return DEFAULT_SYS_PROMPTS[self.task]["empty"][tar_lang] + return DEFAULT_SYS_PROMPTS[self.task][tar_lang] + + def __call__(self, + prompt, + system_prompt=None, + tar_lang="zh", + image=None, + seed=-1, + *args, + **kwargs): + if system_prompt is None: + system_prompt = self.decide_system_prompt( + tar_lang=tar_lang, prompt=prompt) + if seed < 0: + seed = random.randint(0, sys.maxsize) + if image is not None and self.is_vl: + return self.extend_with_img( + prompt, system_prompt, image=image, seed=seed, *args, **kwargs) + elif not self.is_vl: + return self.extend(prompt, system_prompt, seed, *args, **kwargs) + else: + raise NotImplementedError + + +class DashScopePromptExpander(PromptExpander): + + def __init__(self, + api_key=None, + model_name=None, + task=None, + max_image_size=512 * 512, + retry_times=4, + is_vl=False, + **kwargs): + ''' + Args: + api_key: The API key for Dash Scope authentication and access to related services. + model_name: Model name, 'qwen-plus' for extending prompts, 'qwen-vl-max' for extending prompt-images. + task: Task name. This is required to determine the default system prompt. + max_image_size: The maximum size of the image; unit unspecified (e.g., pixels, KB). Please specify the unit based on actual usage. + retry_times: Number of retry attempts in case of request failure. + is_vl: A flag indicating whether the task involves visual-language processing. + **kwargs: Additional keyword arguments that can be passed to the function or method. + ''' + if model_name is None: + model_name = 'qwen-plus' if not is_vl else 'qwen-vl-max' + super().__init__(model_name, task, is_vl, **kwargs) + if api_key is not None: + dashscope.api_key = api_key + elif 'DASH_API_KEY' in os.environ and os.environ[ + 'DASH_API_KEY'] is not None: + dashscope.api_key = os.environ['DASH_API_KEY'] + else: + raise ValueError("DASH_API_KEY is not set") + if 'DASH_API_URL' in os.environ and os.environ[ + 'DASH_API_URL'] is not None: + dashscope.base_http_api_url = os.environ['DASH_API_URL'] + else: + dashscope.base_http_api_url = 'https://dashscope.aliyuncs.com/api/v1' + self.api_key = api_key + + self.max_image_size = max_image_size + self.model = model_name + self.retry_times = retry_times + + def extend(self, prompt, system_prompt, seed=-1, *args, **kwargs): + messages = [{ + 'role': 'system', + 'content': system_prompt + }, { + 'role': 'user', + 'content': prompt + }] + + exception = None + for _ in range(self.retry_times): + try: + response = dashscope.Generation.call( + self.model, + messages=messages, + seed=seed, + result_format='message', # set the result to be "message" format. + ) + assert response.status_code == HTTPStatus.OK, response + expanded_prompt = response['output']['choices'][0]['message'][ + 'content'] + return PromptOutput( + status=True, + prompt=expanded_prompt, + seed=seed, + system_prompt=system_prompt, + message=json.dumps(response, ensure_ascii=False)) + except Exception as e: + exception = e + return PromptOutput( + status=False, + prompt=prompt, + seed=seed, + system_prompt=system_prompt, + message=str(exception)) + + def extend_with_img(self, + prompt, + system_prompt, + image: Union[Image.Image, str] = None, + seed=-1, + *args, + **kwargs): + if isinstance(image, str): + image = Image.open(image).convert('RGB') + w = image.width + h = image.height + area = min(w * h, self.max_image_size) + aspect_ratio = h / w + resized_h = round(math.sqrt(area * aspect_ratio)) + resized_w = round(math.sqrt(area / aspect_ratio)) + image = image.resize((resized_w, resized_h)) + with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as f: + image.save(f.name) + fname = f.name + image_path = f"file://{f.name}" + prompt = f"{prompt}" + messages = [ + { + 'role': 'system', + 'content': [{ + "text": system_prompt + }] + }, + { + 'role': 'user', + 'content': [{ + "text": prompt + }, { + "image": image_path + }] + }, + ] + response = None + result_prompt = prompt + exception = None + status = False + for _ in range(self.retry_times): + try: + response = dashscope.MultiModalConversation.call( + self.model, + messages=messages, + seed=seed, + result_format='message', # set the result to be "message" format. + ) + assert response.status_code == HTTPStatus.OK, response + result_prompt = response['output']['choices'][0]['message'][ + 'content'][0]['text'].replace('\n', '\\n') + status = True + break + except Exception as e: + exception = e + result_prompt = result_prompt.replace('\n', '\\n') + os.remove(fname) + + return PromptOutput( + status=status, + prompt=result_prompt, + seed=seed, + system_prompt=system_prompt, + message=str(exception) if not status else json.dumps( + response, ensure_ascii=False)) + + +class QwenPromptExpander(PromptExpander): + model_dict = { + "QwenVL2.5_3B": "Qwen/Qwen2.5-VL-3B-Instruct", + "QwenVL2.5_7B": "Qwen/Qwen2.5-VL-7B-Instruct", + "Qwen2.5_3B": "Qwen/Qwen2.5-3B-Instruct", + "Qwen2.5_7B": "Qwen/Qwen2.5-7B-Instruct", + "Qwen2.5_14B": "Qwen/Qwen2.5-14B-Instruct", + } + + def __init__(self, + model_name=None, + task=None, + device=0, + is_vl=False, + **kwargs): + ''' + Args: + model_name: Use predefined model names such as 'QwenVL2.5_7B' and 'Qwen2.5_14B', + which are specific versions of the Qwen model. Alternatively, you can use the + local path to a downloaded model or the model name from Hugging Face." + Detailed Breakdown: + Predefined Model Names: + * 'QwenVL2.5_7B' and 'Qwen2.5_14B' are specific versions of the Qwen model. + Local Path: + * You can provide the path to a model that you have downloaded locally. + Hugging Face Model Name: + * You can also specify the model name from Hugging Face's model hub. + task: Task name. This is required to determine the default system prompt. + is_vl: A flag indicating whether the task involves visual-language processing. + **kwargs: Additional keyword arguments that can be passed to the function or method. + ''' + if model_name is None: + model_name = 'Qwen2.5_14B' if not is_vl else 'QwenVL2.5_7B' + super().__init__(model_name, task, is_vl, device, **kwargs) + if (not os.path.exists(self.model_name)) and (self.model_name + in self.model_dict): + self.model_name = self.model_dict[self.model_name] + + if self.is_vl: + # default: Load the model on the available device(s) + from transformers import ( + AutoProcessor, + AutoTokenizer, + Qwen2_5_VLForConditionalGeneration, + ) + try: + from .qwen_vl_utils import process_vision_info + except: + from qwen_vl_utils import process_vision_info + self.process_vision_info = process_vision_info + min_pixels = 256 * 28 * 28 + max_pixels = 1280 * 28 * 28 + self.processor = AutoProcessor.from_pretrained( + self.model_name, + min_pixels=min_pixels, + max_pixels=max_pixels, + use_fast=True) + self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained( + self.model_name, + torch_dtype=torch.bfloat16 if FLASH_VER == 2 else + torch.float16 if "AWQ" in self.model_name else "auto", + attn_implementation="flash_attention_2" + if FLASH_VER == 2 else None, + device_map="cpu") + else: + from transformers import AutoModelForCausalLM, AutoTokenizer + self.model = AutoModelForCausalLM.from_pretrained( + self.model_name, + torch_dtype=torch.float16 + if "AWQ" in self.model_name else "auto", + attn_implementation="flash_attention_2" + if FLASH_VER == 2 else None, + device_map="cpu") + self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) + + def extend(self, prompt, system_prompt, seed=-1, *args, **kwargs): + self.model = self.model.to(self.device) + messages = [{ + "role": "system", + "content": system_prompt + }, { + "role": "user", + "content": prompt + }] + text = self.tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True) + model_inputs = self.tokenizer([text], + return_tensors="pt").to(self.model.device) + + generated_ids = self.model.generate(**model_inputs, max_new_tokens=512) + generated_ids = [ + output_ids[len(input_ids):] for input_ids, output_ids in zip( + model_inputs.input_ids, generated_ids) + ] + + expanded_prompt = self.tokenizer.batch_decode( + generated_ids, skip_special_tokens=True)[0] + self.model = self.model.to("cpu") + return PromptOutput( + status=True, + prompt=expanded_prompt, + seed=seed, + system_prompt=system_prompt, + message=json.dumps({"content": expanded_prompt}, + ensure_ascii=False)) + + def extend_with_img(self, + prompt, + system_prompt, + image: Union[Image.Image, str] = None, + seed=-1, + *args, + **kwargs): + self.model = self.model.to(self.device) + messages = [{ + 'role': 'system', + 'content': [{ + "type": "text", + "text": system_prompt + }] + }, { + "role": + "user", + "content": [ + { + "type": "image", + "image": image, + }, + { + "type": "text", + "text": prompt + }, + ], + }] + + # Preparation for inference + text = self.processor.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True) + image_inputs, video_inputs = self.process_vision_info(messages) + inputs = self.processor( + text=[text], + images=image_inputs, + videos=video_inputs, + padding=True, + return_tensors="pt", + ) + inputs = inputs.to(self.device) + + # Inference: Generation of the output + generated_ids = self.model.generate(**inputs, max_new_tokens=512) + generated_ids_trimmed = [ + out_ids[len(in_ids):] + for in_ids, out_ids in zip(inputs.input_ids, generated_ids) + ] + expanded_prompt = self.processor.batch_decode( + generated_ids_trimmed, + skip_special_tokens=True, + clean_up_tokenization_spaces=False)[0] + self.model = self.model.to("cpu") + return PromptOutput( + status=True, + prompt=expanded_prompt, + seed=seed, + system_prompt=system_prompt, + message=json.dumps({"content": expanded_prompt}, + ensure_ascii=False)) + + +if __name__ == "__main__": + logging.basicConfig( + level=logging.INFO, + format="[%(asctime)s] %(levelname)s: %(message)s", + handlers=[logging.StreamHandler(stream=sys.stdout)]) + + seed = 100 + prompt = "夏日海滩度假风格，一只戴着墨镜的白色猫咪坐在冲浪板上。猫咪毛发蓬松，表情悠闲，直视镜头。背景是模糊的海滩景色，海水清澈，远处有绿色的山丘和蓝天白云。猫咪的姿态自然放松，仿佛在享受海风和阳光。近景特写，强调猫咪的细节和海滩的清新氛围。" + en_prompt = "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." + image = "./examples/i2v_input.JPG" + + def test(method, + prompt, + model_name, + task, + image=None, + en_prompt=None, + seed=None): + prompt_expander = method( + model_name=model_name, task=task, is_vl=image is not None) + result = prompt_expander(prompt, image=image, tar_lang="zh") + logging.info(f"zh prompt -> zh: {result.prompt}") + result = prompt_expander(prompt, image=image, tar_lang="en") + logging.info(f"zh prompt -> en: {result.prompt}") + if en_prompt is not None: + result = prompt_expander(en_prompt, image=image, tar_lang="zh") + logging.info(f"en prompt -> zh: {result.prompt}") + result = prompt_expander(en_prompt, image=image, tar_lang="en") + logging.info(f"en prompt -> en: {result.prompt}") + + ds_model_name = None + ds_vl_model_name = None + qwen_model_name = None + qwen_vl_model_name = None + + for task in ["t2v-A14B", "i2v-A14B", "ti2v-5B"]: + # test prompt extend + if "t2v" in task or "ti2v" in task: + # test dashscope api + logging.info(f"-" * 40) + logging.info(f"Testing {task} dashscope prompt extend") + test( + DashScopePromptExpander, + prompt, + ds_model_name, + task, + image=None, + en_prompt=en_prompt, + seed=seed) + + # test qwen api + logging.info(f"-" * 40) + logging.info(f"Testing {task} qwen prompt extend") + test( + QwenPromptExpander, + prompt, + qwen_model_name, + task, + image=None, + en_prompt=en_prompt, + seed=seed) + + # test prompt-image extend + if "i2v" in task: + # test dashscope api + logging.info(f"-" * 40) + logging.info(f"Testing {task} dashscope vl prompt extend") + test( + DashScopePromptExpander, + prompt, + ds_vl_model_name, + task, + image=image, + en_prompt=en_prompt, + seed=seed) + + # test qwen api + logging.info(f"-" * 40) + logging.info(f"Testing {task} qwen vl prompt extend") + test( + QwenPromptExpander, + prompt, + qwen_vl_model_name, + task, + image=image, + en_prompt=en_prompt, + seed=seed) + + # test empty prompt extend + if "i2v-A14B" in task: + # test dashscope api + logging.info(f"-" * 40) + logging.info(f"Testing {task} dashscope vl empty prompt extend") + test( + DashScopePromptExpander, + "", + ds_vl_model_name, + task, + image=image, + en_prompt=None, + seed=seed) + + # test qwen api + logging.info(f"-" * 40) + logging.info(f"Testing {task} qwen vl empty prompt extend") + test( + QwenPromptExpander, + "", + qwen_vl_model_name, + task, + image=image, + en_prompt=None, + seed=seed) diff --git a/wan_5b/utils/qwen_vl_utils.py b/wan_5b/utils/qwen_vl_utils.py new file mode 100644 index 0000000..bf0e832 --- /dev/null +++ b/wan_5b/utils/qwen_vl_utils.py @@ -0,0 +1,363 @@ +# Copied from https://github.com/kq-chen/qwen-vl-utils +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +from __future__ import annotations + +import base64 +import logging +import math +import os +import sys +import time +import warnings +from functools import lru_cache +from io import BytesIO + +import requests +import torch +import torchvision +from packaging import version +from PIL import Image +from torchvision import io, transforms +from torchvision.transforms import InterpolationMode + +logger = logging.getLogger(__name__) + +IMAGE_FACTOR = 28 +MIN_PIXELS = 4 * 28 * 28 +MAX_PIXELS = 16384 * 28 * 28 +MAX_RATIO = 200 + +VIDEO_MIN_PIXELS = 128 * 28 * 28 +VIDEO_MAX_PIXELS = 768 * 28 * 28 +VIDEO_TOTAL_PIXELS = 24576 * 28 * 28 +FRAME_FACTOR = 2 +FPS = 2.0 +FPS_MIN_FRAMES = 4 +FPS_MAX_FRAMES = 768 + + +def round_by_factor(number: int, factor: int) -> int: + """Returns the closest integer to 'number' that is divisible by 'factor'.""" + return round(number / factor) * factor + + +def ceil_by_factor(number: int, factor: int) -> int: + """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'.""" + return math.ceil(number / factor) * factor + + +def floor_by_factor(number: int, factor: int) -> int: + """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'.""" + return math.floor(number / factor) * factor + + +def smart_resize(height: int, + width: int, + factor: int = IMAGE_FACTOR, + min_pixels: int = MIN_PIXELS, + max_pixels: int = MAX_PIXELS) -> tuple[int, int]: + """ + Rescales the image so that the following conditions are met: + + 1. Both dimensions (height and width) are divisible by 'factor'. + + 2. The total number of pixels is within the range ['min_pixels', 'max_pixels']. + + 3. The aspect ratio of the image is maintained as closely as possible. + """ + if max(height, width) / min(height, width) > MAX_RATIO: + raise ValueError( + f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}" + ) + h_bar = max(factor, round_by_factor(height, factor)) + w_bar = max(factor, round_by_factor(width, factor)) + if h_bar * w_bar > max_pixels: + beta = math.sqrt((height * width) / max_pixels) + h_bar = floor_by_factor(height / beta, factor) + w_bar = floor_by_factor(width / beta, factor) + elif h_bar * w_bar < min_pixels: + beta = math.sqrt(min_pixels / (height * width)) + h_bar = ceil_by_factor(height * beta, factor) + w_bar = ceil_by_factor(width * beta, factor) + return h_bar, w_bar + + +def fetch_image(ele: dict[str, str | Image.Image], + size_factor: int = IMAGE_FACTOR) -> Image.Image: + if "image" in ele: + image = ele["image"] + else: + image = ele["image_url"] + image_obj = None + if isinstance(image, Image.Image): + image_obj = image + elif image.startswith("http://") or image.startswith("https://"): + image_obj = Image.open(requests.get(image, stream=True).raw) + elif image.startswith("file://"): + image_obj = Image.open(image[7:]) + elif image.startswith("data:image"): + if "base64," in image: + _, base64_data = image.split("base64,", 1) + data = base64.b64decode(base64_data) + image_obj = Image.open(BytesIO(data)) + else: + image_obj = Image.open(image) + if image_obj is None: + raise ValueError( + f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}" + ) + image = image_obj.convert("RGB") + ## resize + if "resized_height" in ele and "resized_width" in ele: + resized_height, resized_width = smart_resize( + ele["resized_height"], + ele["resized_width"], + factor=size_factor, + ) + else: + width, height = image.size + min_pixels = ele.get("min_pixels", MIN_PIXELS) + max_pixels = ele.get("max_pixels", MAX_PIXELS) + resized_height, resized_width = smart_resize( + height, + width, + factor=size_factor, + min_pixels=min_pixels, + max_pixels=max_pixels, + ) + image = image.resize((resized_width, resized_height)) + + return image + + +def smart_nframes( + ele: dict, + total_frames: int, + video_fps: int | float, +) -> int: + """calculate the number of frames for video used for model inputs. + + Args: + ele (dict): a dict contains the configuration of video. + support either `fps` or `nframes`: + - nframes: the number of frames to extract for model inputs. + - fps: the fps to extract frames for model inputs. + - min_frames: the minimum number of frames of the video, only used when fps is provided. + - max_frames: the maximum number of frames of the video, only used when fps is provided. + total_frames (int): the original total number of frames of the video. + video_fps (int | float): the original fps of the video. + + Raises: + ValueError: nframes should in interval [FRAME_FACTOR, total_frames]. + + Returns: + int: the number of frames for video used for model inputs. + """ + assert not ("fps" in ele and + "nframes" in ele), "Only accept either `fps` or `nframes`" + if "nframes" in ele: + nframes = round_by_factor(ele["nframes"], FRAME_FACTOR) + else: + fps = ele.get("fps", FPS) + min_frames = ceil_by_factor( + ele.get("min_frames", FPS_MIN_FRAMES), FRAME_FACTOR) + max_frames = floor_by_factor( + ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)), + FRAME_FACTOR) + nframes = total_frames / video_fps * fps + nframes = min(max(nframes, min_frames), max_frames) + nframes = round_by_factor(nframes, FRAME_FACTOR) + if not (FRAME_FACTOR <= nframes and nframes <= total_frames): + raise ValueError( + f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}." + ) + return nframes + + +def _read_video_torchvision(ele: dict,) -> torch.Tensor: + """read video using torchvision.io.read_video + + Args: + ele (dict): a dict contains the configuration of video. + support keys: + - video: the path of video. support "file://", "http://", "https://" and local path. + - video_start: the start time of video. + - video_end: the end time of video. + Returns: + torch.Tensor: the video tensor with shape (T, C, H, W). + """ + video_path = ele["video"] + if version.parse(torchvision.__version__) < version.parse("0.19.0"): + if "http://" in video_path or "https://" in video_path: + warnings.warn( + "torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0." + ) + if "file://" in video_path: + video_path = video_path[7:] + st = time.time() + video, audio, info = io.read_video( + video_path, + start_pts=ele.get("video_start", 0.0), + end_pts=ele.get("video_end", None), + pts_unit="sec", + output_format="TCHW", + ) + total_frames, video_fps = video.size(0), info["video_fps"] + logger.info( + f"torchvision: {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s" + ) + nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps) + idx = torch.linspace(0, total_frames - 1, nframes).round().long() + video = video[idx] + return video + + +def is_decord_available() -> bool: + import importlib.util + + return importlib.util.find_spec("decord") is not None + + +def _read_video_decord(ele: dict,) -> torch.Tensor: + """read video using decord.VideoReader + + Args: + ele (dict): a dict contains the configuration of video. + support keys: + - video: the path of video. support "file://", "http://", "https://" and local path. + - video_start: the start time of video. + - video_end: the end time of video. + Returns: + torch.Tensor: the video tensor with shape (T, C, H, W). + """ + import decord + video_path = ele["video"] + st = time.time() + vr = decord.VideoReader(video_path) + # TODO: support start_pts and end_pts + if 'video_start' in ele or 'video_end' in ele: + raise NotImplementedError( + "not support start_pts and end_pts in decord for now.") + total_frames, video_fps = len(vr), vr.get_avg_fps() + logger.info( + f"decord: {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s" + ) + nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps) + idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist() + video = vr.get_batch(idx).asnumpy() + video = torch.tensor(video).permute(0, 3, 1, 2) # Convert to TCHW format + return video + + +VIDEO_READER_BACKENDS = { + "decord": _read_video_decord, + "torchvision": _read_video_torchvision, +} + +FORCE_QWENVL_VIDEO_READER = os.getenv("FORCE_QWENVL_VIDEO_READER", None) + + +@lru_cache(maxsize=1) +def get_video_reader_backend() -> str: + if FORCE_QWENVL_VIDEO_READER is not None: + video_reader_backend = FORCE_QWENVL_VIDEO_READER + elif is_decord_available(): + video_reader_backend = "decord" + else: + video_reader_backend = "torchvision" + logger.info( + f"qwen-vl-utils using {video_reader_backend} to read video.", + file=sys.stderr) + return video_reader_backend + + +def fetch_video( + ele: dict, + image_factor: int = IMAGE_FACTOR) -> torch.Tensor | list[Image.Image]: + if isinstance(ele["video"], str): + video_reader_backend = get_video_reader_backend() + video = VIDEO_READER_BACKENDS[video_reader_backend](ele) + nframes, _, height, width = video.shape + + min_pixels = ele.get("min_pixels", VIDEO_MIN_PIXELS) + total_pixels = ele.get("total_pixels", VIDEO_TOTAL_PIXELS) + max_pixels = max( + min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR), + int(min_pixels * 1.05)) + max_pixels = ele.get("max_pixels", max_pixels) + if "resized_height" in ele and "resized_width" in ele: + resized_height, resized_width = smart_resize( + ele["resized_height"], + ele["resized_width"], + factor=image_factor, + ) + else: + resized_height, resized_width = smart_resize( + height, + width, + factor=image_factor, + min_pixels=min_pixels, + max_pixels=max_pixels, + ) + video = transforms.functional.resize( + video, + [resized_height, resized_width], + interpolation=InterpolationMode.BICUBIC, + antialias=True, + ).float() + return video + else: + assert isinstance(ele["video"], (list, tuple)) + process_info = ele.copy() + process_info.pop("type", None) + process_info.pop("video", None) + images = [ + fetch_image({ + "image": video_element, + **process_info + }, + size_factor=image_factor) + for video_element in ele["video"] + ] + nframes = ceil_by_factor(len(images), FRAME_FACTOR) + if len(images) < nframes: + images.extend([images[-1]] * (nframes - len(images))) + return images + + +def extract_vision_info( + conversations: list[dict] | list[list[dict]]) -> list[dict]: + vision_infos = [] + if isinstance(conversations[0], dict): + conversations = [conversations] + for conversation in conversations: + for message in conversation: + if isinstance(message["content"], list): + for ele in message["content"]: + if ("image" in ele or "image_url" in ele or + "video" in ele or + ele["type"] in ("image", "image_url", "video")): + vision_infos.append(ele) + return vision_infos + + +def process_vision_info( + conversations: list[dict] | list[list[dict]], +) -> tuple[list[Image.Image] | None, list[torch.Tensor | list[Image.Image]] | + None]: + vision_infos = extract_vision_info(conversations) + ## Read images or videos + image_inputs = [] + video_inputs = [] + for vision_info in vision_infos: + if "image" in vision_info or "image_url" in vision_info: + image_inputs.append(fetch_image(vision_info)) + elif "video" in vision_info: + video_inputs.append(fetch_video(vision_info)) + else: + raise ValueError("image, image_url or video should in content.") + if len(image_inputs) == 0: + image_inputs = None + if len(video_inputs) == 0: + video_inputs = None + return image_inputs, video_inputs diff --git a/wan_5b/utils/system_prompt.py b/wan_5b/utils/system_prompt.py new file mode 100644 index 0000000..c494705 --- /dev/null +++ b/wan_5b/utils/system_prompt.py @@ -0,0 +1,147 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. + +T2V_A14B_ZH_SYS_PROMPT = \ +''' 你是一位电影导演，旨在为用户输入的原始prompt添加电影元素，改写为优质Prompt，使其完整、具有表现力。 +任务要求： +1. 对于用户输入的prompt,在不改变prompt的原意（如主体、动作）前提下，从下列电影美学设定中选择部分合适的时间、光源、光线强度、光线角度、对比度、饱和度、色调、拍摄角度、镜头大小、构图的电影设定细节,将这些内容添加到prompt中，让画面变得更美，注意，可以任选，不必每项都有 + 时间：["白天", "夜晚", "黎明", "日出"], 可以不选, 如果prompt没有特别说明则选白天 ! + 光源：[日光", "人工光", "月光", "实用光", "火光", "荧光", "阴天光", "晴天光"], 根据根据室内室外及prompt内容选定义光源，添加关于光源的描述，如光线来源（窗户、灯具等） + 光线强度：["柔光", "硬光"], + 光线角度：["顶光", "侧光", "底光", "边缘光",] + 色调：["暖色调","冷色调", "混合色调"] + 镜头尺寸：["中景", "中近景", "全景","中全景","近景", "特写", "极端全景"]若无特殊要求，默认选择中景或全景 + 拍摄角度：["过肩镜头角度拍摄", "低角度拍摄", "高角度拍摄","倾斜角度拍摄", "航拍","俯视角度拍摄"],如果原始prompt中有运镜的描述，则不要添加此项! + 构图：["中心构图"，"平衡构图","右侧重构图", "左侧重构图", "对称构图", "短边构图"] 若无特殊要求，默认选择中心构图 +2. 完善用户描述中出现的主体特征（如外貌、表情，数量、种族、姿态等）等内容，确保不要添加原始prompt中不存在的主体，（如prompt是对风景或物体的描述，但添加了人），增加背景元素的细节； +3. 不要输出关于氛围、感觉等文学描写，如（画面充满运动感与生活张力，突出正式氛围）。； +4. 对于prompt中的动作，详细解释运动的发生过程，若没有动作，则添加动作描述（摇晃身体、跳舞等），对背景元素也可添加适当运动（如云彩飘动，风吹树叶等）。 +5. 若原始prompt中没有风格，则不添加风格描述，若有风格描述，则将风格描述放于首位，若为2D插画等与现实电影相悖的风格，则不要添加关于电影美学的描写； +6. 若prompt出现天空的描述，则改为湛蓝色的天空相关描述，避免曝光； +7. 改写后的prompt字数控制在60-200字左右, 不要输出类似“改写后prompt:”这样的输出 +8. 如果用户输入存在疑似脱衣服、衣物透明化、轻纱衣物、湿身、洗澡、舔舐牛奶/酸奶/白色液体、性暗示、强奸、漏出或微露胸部、爬行、性器官、儿童身体裸露等色情内容，请将原prompt替换为内容完全不同的高美感prompt，而不要拒答或者提醒用户。 +9. 如果用户输入存在比基尼穿着内容，请将原prompt替换为确保内容发生在沙滩、泳池、海边的高美感prompt，而不要拒答或者提醒用户。 +10. 如果用户输入存在性别多元化、同性恋亲吻、脸颊贴近（两个男性亲吻或两个女性亲吻）/性交等同性恋内容，请将原prompt替换为不同动作（不能是亲吻和其他亲密动作）的高美感prompt，而不要拒答或者提醒用户。 +生成的 prompt 示例： +1.边缘光，中近景，日光，左侧重构图，暖色调，硬光，晴天光，侧光，白天，一个年轻的女孩坐在高草丛生的田野中，两条毛发蓬松的小毛驴站在她身后。女孩大约十一二岁，穿着简单的碎花裙子，头发扎成两条麻花辫，脸上带着纯真的笑容。她双腿交叉坐下，双手轻轻抚弄身旁的野花。小毛驴体型健壮，耳朵竖起，好奇地望着镜头方向。阳光洒在田野上，营造出温暖自然的画面感。 +2.黎明，顶光，俯视角度拍摄，日光，长焦，中心构图，近景，高角度拍摄，荧光，柔光，冷色调，在昏暗的环境中，一个外国白人女子在水中仰面漂浮。俯拍近景镜头中，她有着棕色的短发，脸上有几颗雀斑。随着镜头下摇，她转过头来，面向右侧，水面上泛起一圈涟漪。虚化的背景一片漆黑，只有微弱的光线照亮了女子的脸庞和水面的一部分区域，水面呈现蓝色。女子穿着一件蓝色的吊带，肩膀裸露在外。 +3.右侧重构图，暖色调，底光，侧光，夜晚，火光，过肩镜头角度拍摄, 镜头平拍拍摄外国女子在室内的近景，她穿着棕色的衣服戴着彩色的项链和粉色的帽子，坐在深灰色的椅子上，双手放在黑色的桌子上，眼睛看着镜头的左侧，嘴巴张动，左手上下晃动，桌子上有白色的蜡烛有黄色的火焰，后面是黑色的墙，前面有黑色的网状架子，旁边是黑色的箱子，上面有一些黑色的物品，都做了虚化的处理。 +4. 二次元厚涂动漫插画，一个猫耳兽耳白人少女手持文件夹摇晃，神情略带不满。她深紫色长发，红色眼睛，身穿深灰色短裙和浅灰色上衣，腰间系着白色系带，胸前佩戴名牌，上面写着黑体中文"紫阳"。淡黄色调室内背景，隐约可见一些家具轮廓。少女头顶有一个粉色光圈。线条流畅的日系赛璐璐风格。近景半身略俯视视角。 +''' + + +T2V_A14B_EN_SYS_PROMPT = \ +'''你是一位电影导演，旨在为用户输入的原始prompt添加电影元素，改写为优质（英文）Prompt，使其完整、具有表现力注意，输出必须是英文！ +任务要求： +1. 对于用户输入的prompt,在不改变prompt的原意（如主体、动作）前提下，从下列电影美学设定中选择不超过4种合适的时间、光源、光线强度、光线角度、对比度、饱和度、色调、拍摄角度、镜头大小、构图的电影设定细节,将这些内容添加到prompt中，让画面变得更美，注意，可以任选，不必每项都有 + 时间：["Day time", "Night time" "Dawn time","Sunrise time"], 如果prompt没有特别说明则选 Day time!!! + 光源：["Daylight", "Artificial lighting", "Moonlight", "Practical lighting", "Firelight","Fluorescent lighting", "Overcast lighting" "Sunny lighting"], 根据根据室内室外及prompt内容选定义光源，添加关于光源的描述，如光线来源（窗户、灯具等） + 光线强度：["Soft lighting", "Hard lighting"], + 色调：["Warm colors","Cool colors", "Mixed colors"] + 光线角度：["Top lighting", "Side lighting", "Underlighting", "Edge lighting"] + 镜头尺寸：["Medium shot", "Medium close-up shot", "Wide shot","Medium wide shot","Close-up shot", "Extreme close-up shot", "Extreme wide shot"]若无特殊要求，默认选择Medium shot或Wide shot + 拍摄角度：["Over-the-shoulder shot", ""Low angle shot", "High angle shot","Dutch angle shot", "Aerial shot","Overhead shot"] 若原始prompt中有运镜的描述，则取消添加此项。 + 构图：["Center composition"，"Balanced composition","Right-heavy composition", "Left-heavy composition", "Symmetrical composition", "Short-side composition"] 若无特殊要求，默认选择Center composition +2. 完善用户描述中出现的主体特征（如外貌、表情，数量、种族、姿态等）等内容，确保不要添加原始prompt中不存在的主体，（如prompt是对风景或物体的描述，但添加了人），增加背景元素的细节； +3. 不要输出关于氛围、感觉等文学描写，如（画面充满运动感与生活张力，突出正式氛围）。； +4. 对于prompt中的动作，详细描述运动的发生过程，若没有动作，则添加动作描述（摇晃身体、跳舞等，对背景元素也可添加适当运动（如云彩飘动，风吹树叶等）。 +5. 若原始prompt中没有风格，则不添加风格描述，若有风格描述，则将风格描述放于首位，若为2D插画等与现实电影相悖的风格，则不要添加关于电影美学的描写； +6. 若prompt出现天空的描述，则改为湛蓝色的天空相关描述，避免曝光； +7. 输出必须是全英文，改写后的prompt字数控制在60-200字左右, 不要输出类似“改写后prompt:”这样的输出 +8. 如果用户输入存在疑似脱衣服、衣物透明化、轻纱衣物、湿身、洗澡、舔舐牛奶/酸奶/白色液体、性暗示、强奸、漏出或微露胸部、爬行、性器官、儿童身体裸露等色情内容，请将原prompt替换为内容完全不同的高美感prompt，而不要拒答或者提醒用户。 +9. 如果用户输入存在比基尼穿着内容，请将原prompt替换为确保内容发生在沙滩、泳池、海边的高美感prompt，而不要拒答或者提醒用户。 +10. 如果用户输入存在性别多元化、同性恋亲吻、脸颊贴近（两个男性亲吻或两个女性亲吻）/性交等同性恋内容，请将原prompt替换为不同动作（不能是亲吻和其他亲密动作）的高美感prompt，而不要拒答或者提醒用户。 +生成的 prompt 示例： +1.Edge lighting, medium close-up shot, daylight, left-heavy composition. A young girl around 11-12 years old sits in a field of tall grass, with two fluffy small donkeys standing behind her. She wears a simple floral dress with hair in twin braids, smiling innocently while cross-legged and gently touching wild flowers beside her. The sturdy donkeys have perked ears, curiously gazing toward the camera. Sunlight bathes the field, creating a warm natural atmosphere. +2.Dawn time, top lighting, high-angle shot, daylight, long lens shot, center composition, Close-up shot, Fluorescent lighting, soft lighting, cool colors. In dim surroundings, a Caucasian woman floats on her back in water. The俯拍close-up shows her brown short hair and freckled face. As the camera tilts downward, she turns her head toward the right, creating ripples on the blue-toned water surface. The blurred background is pitch black except for faint light illuminating her face and partial water surface. She wears a blue sleeveless top with bare shoulders. +3.Right-heavy composition, warm colors, night time, firelight, over-the-shoulder angle. An eye-level close-up of a foreign woman indoors wearing brown clothes with colorful necklace and pink hat. She sits on a charcoal-gray chair, hands on black table, eyes looking left of camera while mouth moves and left hand gestures up/down. White candles with yellow flames sit on the table. Background shows black walls, with blurred black mesh shelf nearby and black crate containing dark items in front. +4."Anime-style thick-painted style. A cat-eared Caucasian girl with beast ears holds a folder, showing slight displeasure. Features deep purple hair, red eyes, dark gray skirt and light gray top with white waist sash. A name tag labeled 'Ziyang' in bold Chinese characters hangs on her chest. Pale yellow indoor background with faint furniture outlines. A pink halo floats above her head. Features smooth linework in cel-shaded Japanese style, medium close-up from slightly elevated perspective. +''' + + +I2V_A14B_ZH_SYS_PROMPT = \ +'''你是一个视频描述提示词的改写专家，你的任务是根据用户给你输入的图像，对提供的视频描述提示词进行改写，你要强调潜在的动态内容。具体要求如下 +用户输入的语言可能含有多样化的描述，如markdown文档格式、指令格式，长度过长或者过短，你需要根据图片的内容和用户的输入的提示词，尽可能提取用户输入的提示词和图片关联信息。 +你改写的视频描述结果要尽可能保留提供给你的视频描述提示词中动态部分，保留主体的动作。 +你要根据图像，强调并简化视频描述提示词中的图像主体，如果用户只提供了动作，你要根据图像内容合理补充，如“跳舞”补充称“一个女孩在跳舞” +如果用户输入的提示词过长，你需要提炼潜在的动作过程 +如果用户输入的提示词过短，综合用户输入的提示词以及画面内容，合理的增加潜在的运动信息 +你要根据图像，保留并强调视频描述提示词中关于运镜手段的描述，如“镜头上摇”，“镜头从左到右”，“镜头从右到左”等等，你要保留，如“镜头拍摄两个男人打斗，他们先是躺在地上，随后镜头向上移动，拍摄他们站起来，接着镜头向左移动，左边男人拿着一个蓝色的东西，右边男人上前抢夺，两人激烈地来回争抢。”。 +你需要给出对视频描述的动态内容，不要添加对于静态场景的描述，如果用户输入的描述已经在画面中出现，则移除这些描述 +改写后的prompt字数控制在100字以下 +无论用户输入那种语言，你都需要输出中文 +改写后 prompt 示例： +1. 镜头后拉，拍摄两个外国男人，走在楼梯上，镜头左侧的男人右手搀扶着镜头右侧的男人。 +2. 一只黑色的小松鼠专注地吃着东西，偶尔抬头看看四周。 +3. 男子说着话，表情从微笑逐渐转变为闭眼，然后睁开眼睛，最后是闭眼微笑，他的手势活跃，在说话时做出一系列的手势。 +4. 一个人正在用尺子和笔进行测量的特写，右手用一支黑色水性笔在纸上画出一条直线。 +5. 一辆车模型在木板上形式，车辆从画面的右侧向左侧移动，经过一片草地和一些木制结构。 +6. 镜头左移后前推，拍摄一个人坐在防波堤上。 +7. 男子说着话，他的表情和手势随着对话内容的变化而变化，但整体场景保持不变。 +8. 镜头左移后前推，拍摄一个人坐在防波堤上。 +9. 带着珍珠项链的女子看向画面右侧并说着话。 +请直接输出改写后的文本，不要进行多余的回复。''' + + +I2V_A14B_EN_SYS_PROMPT = \ +'''You are an expert in rewriting video description prompts. Your task is to rewrite the provided video description prompts based on the images given by users, emphasizing potential dynamic content. Specific requirements are as follows: +The user's input language may include diverse descriptions, such as markdown format, instruction format, or be too long or too short. You need to extract the relevant information from the user’s input and associate it with the image content. +Your rewritten video description should retain the dynamic parts of the provided prompts, focusing on the main subject's actions. Emphasize and simplify the main subject of the image while retaining their movement. If the user only provides an action (e.g., "dancing"), supplement it reasonably based on the image content (e.g., "a girl is dancing"). +If the user’s input prompt is too long, refine it to capture the essential action process. If the input is too short, add reasonable motion-related details based on the image content. +Retain and emphasize descriptions of camera movements, such as "the camera pans up," "the camera moves from left to right," or "the camera moves from right to left." For example: "The camera captures two men fighting. They start lying on the ground, then the camera moves upward as they stand up. The camera shifts left, showing the man on the left holding a blue object while the man on the right tries to grab it, resulting in a fierce back-and-forth struggle." +Focus on dynamic content in the video description and avoid adding static scene descriptions. If the user’s input already describes elements visible in the image, remove those static descriptions. +Limit the rewritten prompt to 100 words or less. Regardless of the input language, your output must be in English. + +Examples of rewritten prompts: +The camera pulls back to show two foreign men walking up the stairs. The man on the left supports the man on the right with his right hand. +A black squirrel focuses on eating, occasionally looking around. +A man talks, his expression shifting from smiling to closing his eyes, reopening them, and finally smiling with closed eyes. His gestures are lively, making various hand motions while speaking. +A close-up of someone measuring with a ruler and pen, drawing a straight line on paper with a black marker in their right hand. +A model car moves on a wooden board, traveling from right to left across grass and wooden structures. +The camera moves left, then pushes forward to capture a person sitting on a breakwater. +A man speaks, his expressions and gestures changing with the conversation, while the overall scene remains constant. +The camera moves left, then pushes forward to capture a person sitting on a breakwater. +A woman wearing a pearl necklace looks to the right and speaks. +Output only the rewritten text without additional responses.''' + + +I2V_A14B_EMPTY_ZH_SYS_PROMPT = \ +'''你是一个视频描述提示词的撰写专家，你的任务是根据用户给你输入的图像，发挥合理的想象，让这张图动起来，你要强调潜在的动态内容。具体要求如下 +你需要根据图片的内容想象出运动的主体 +你输出的结果应强调图片中的动态部分，保留主体的动作。 +你需要给出对视频描述的动态内容，不要有过多的对于静态场景的描述 +输出的prompt字数控制在100字以下 +你需要输出中文 +prompt 示例： +1. 镜头后拉，拍摄两个外国男人，走在楼梯上，镜头左侧的男人右手搀扶着镜头右侧的男人。 +2. 一只黑色的小松鼠专注地吃着东西，偶尔抬头看看四周。 +3. 男子说着话，表情从微笑逐渐转变为闭眼，然后睁开眼睛，最后是闭眼微笑，他的手势活跃，在说话时做出一系列的手势。 +4. 一个人正在用尺子和笔进行测量的特写，右手用一支黑色水性笔在纸上画出一条直线。 +5. 一辆车模型在木板上形式，车辆从画面的右侧向左侧移动，经过一片草地和一些木制结构。 +6. 镜头左移后前推，拍摄一个人坐在防波堤上。 +7. 男子说着话，他的表情和手势随着对话内容的变化而变化，但整体场景保持不变。 +8. 镜头左移后前推，拍摄一个人坐在防波堤上。 +9. 带着珍珠项链的女子看向画面右侧并说着话。 +请直接输出文本，不要进行多余的回复。''' + + +I2V_A14B_EMPTY_EN_SYS_PROMPT = \ +'''You are an expert in writing video description prompts. Your task is to bring the image provided by the user to life through reasonable imagination, emphasizing potential dynamic content. Specific requirements are as follows: + +You need to imagine the moving subject based on the content of the image. +Your output should emphasize the dynamic parts of the image and retain the main subject’s actions. +Focus only on describing dynamic content; avoid excessive descriptions of static scenes. +Limit the output prompt to 100 words or less. +The output must be in English. + +Prompt examples: + +The camera pulls back to show two foreign men walking up the stairs. The man on the left supports the man on the right with his right hand. +A black squirrel focuses on eating, occasionally looking around. +A man talks, his expression shifting from smiling to closing his eyes, reopening them, and finally smiling with closed eyes. His gestures are lively, making various hand motions while speaking. +A close-up of someone measuring with a ruler and pen, drawing a straight line on paper with a black marker in their right hand. +A model car moves on a wooden board, traveling from right to left across grass and wooden structures. +The camera moves left, then pushes forward to capture a person sitting on a breakwater. +A man speaks, his expressions and gestures changing with the conversation, while the overall scene remains constant. +The camera moves left, then pushes forward to capture a person sitting on a breakwater. +A woman wearing a pearl necklace looks to the right and speaks. +Output only the text without additional responses.''' diff --git a/wan_5b/utils/utils.py b/wan_5b/utils/utils.py new file mode 100644 index 0000000..c563c69 --- /dev/null +++ b/wan_5b/utils/utils.py @@ -0,0 +1,159 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import argparse +import binascii +import logging +import os +import os.path as osp + +import imageio +import torch +import torchvision + +__all__ = ['save_video', 'save_image', 'str2bool'] + + +def rand_name(length=8, suffix=''): + name = binascii.b2a_hex(os.urandom(length)).decode('utf-8') + if suffix: + if not suffix.startswith('.'): + suffix = '.' + suffix + name += suffix + return name + + +def save_video(tensor, + save_file=None, + fps=30, + suffix='.mp4', + nrow=8, + normalize=True, + value_range=(-1, 1)): + # cache file + cache_file = osp.join('/tmp', rand_name( + suffix=suffix)) if save_file is None else save_file + + # save to cache + try: + # preprocess + tensor = tensor.clamp(min(value_range), max(value_range)) + tensor = torch.stack([ + torchvision.utils.make_grid( + u, nrow=nrow, normalize=normalize, value_range=value_range) + for u in tensor.unbind(2) + ], + dim=1).permute(1, 2, 3, 0) + tensor = (tensor * 255).type(torch.uint8).cpu() + + # write video + writer = imageio.get_writer( + cache_file, fps=fps, codec='libx264', quality=8) + for frame in tensor.numpy(): + writer.append_data(frame) + writer.close() + except Exception as e: + logging.info(f'save_video failed, error: {e}') + + +def save_image(tensor, save_file, nrow=8, normalize=True, value_range=(-1, 1)): + # cache file + suffix = osp.splitext(save_file)[1] + if suffix.lower() not in [ + '.jpg', '.jpeg', '.png', '.tiff', '.gif', '.webp' + ]: + suffix = '.png' + + # save to cache + try: + tensor = tensor.clamp(min(value_range), max(value_range)) + torchvision.utils.save_image( + tensor, + save_file, + nrow=nrow, + normalize=normalize, + value_range=value_range) + return save_file + except Exception as e: + logging.info(f'save_image failed, error: {e}') + + +def str2bool(v): + """ + Convert a string to a boolean. + + Supported true values: 'yes', 'true', 't', 'y', '1' + Supported false values: 'no', 'false', 'f', 'n', '0' + + Args: + v (str): String to convert. + + Returns: + bool: Converted boolean value. + + Raises: + argparse.ArgumentTypeError: If the value cannot be converted to boolean. + """ + if isinstance(v, bool): + return v + v_lower = v.lower() + if v_lower in ('yes', 'true', 't', 'y', '1'): + return True + elif v_lower in ('no', 'false', 'f', 'n', '0'): + return False + else: + raise argparse.ArgumentTypeError('Boolean value expected (True/False)') + + +def masks_like(tensor, zero=False, generator=None, p=0.2): + assert isinstance(tensor, list) + out1 = [torch.ones(u.shape, dtype=u.dtype, device=u.device) for u in tensor] + + out2 = [torch.ones(u.shape, dtype=u.dtype, device=u.device) for u in tensor] + + if zero: + if generator is not None: + for u, v in zip(out1, out2): + random_num = torch.rand( + 1, generator=generator, device=generator.device).item() + if random_num < p: + u[:, 0] = torch.normal( + mean=-3.5, + std=0.5, + size=(1,), + device=u.device, + generator=generator).expand_as(u[:, 0]).exp() + v[:, 0] = torch.zeros_like(v[:, 0]) + else: + u[:, 0] = u[:, 0] + v[:, 0] = v[:, 0] + else: + for u, v in zip(out1, out2): + u[:, 0] = torch.zeros_like(u[:, 0]) + v[:, 0] = torch.zeros_like(v[:, 0]) + + return out1, out2 + + +def best_output_size(w, h, dw, dh, expected_area): + # float output size + ratio = w / h + ow = (expected_area * ratio)**0.5 + oh = expected_area / ow + + # process width first + ow1 = int(ow // dw * dw) + oh1 = int(expected_area / ow1 // dh * dh) + assert ow1 % dw == 0 and oh1 % dh == 0 and ow1 * oh1 <= expected_area + ratio1 = ow1 / oh1 + + # process height first + oh2 = int(oh // dh * dh) + ow2 = int(expected_area / oh2 // dw * dw) + assert oh2 % dh == 0 and ow2 % dw == 0 and ow2 * oh2 <= expected_area + ratio2 = ow2 / oh2 + + # compare ratios + if max(ratio / ratio1, ratio1 / ratio) < max(ratio / ratio2, + ratio2 / ratio): + return ow1, oh1 + else: + return ow2, oh2