[Feat] Support LTX2 sequence parallel, tensor parallel, and VAE decod… #810

gushiqiao · 2026-01-23T06:18:10Z

…e tiling

gemini-code-assist · 2026-01-23T06:19:01Z

Summary of Changes

Hello @gushiqiao, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces significant architectural enhancements to the LightX2V framework, primarily focusing on improving scalability and performance through the implementation of Tensor Parallelism and Sequence Parallelism. It also streamlines the VAE decoding process with tiling capabilities and refines the image conditioning interface for better usability. These changes aim to make the system more robust and efficient for handling large-scale video generation tasks.

Highlights

Distributed Parallelism (Tensor Parallel & Sequence Parallel): Introduced dedicated classes (MMWeightTP, RMSWeightTP, LTX2AttentionTP, LTX2FFNTP) and logic to support Tensor Parallelism for matrix multiplication, RMS normalization, attention, and FFN layers. Enhanced the LTX2 model to handle data splitting and gathering for both Sequence Parallelism and Tensor Parallelism, including weight distribution from rank 0. Updated parallel configuration settings to correctly configure device meshes for these new modes.
VAE Decode Tiling Support: Added a new configuration option use_tiling_vae and integrated VAE decode tiling into the LTX2VideoVAE, allowing for more memory-efficient video generation by processing VAE decoding in tiles, with automatic application of default tiling settings if enabled.
Improved Image Conditioning API: Refactored the image conditioning arguments from a complex string format (--images) to more intuitive separate arguments (--image_path supporting multiple paths and --image_strength for scalar or per-image strength), simplifying the API for image-to-audio-video (i2av) tasks. Example scripts and input information dataclasses were updated accordingly.
Configuration and Path Validation: Implemented a new validate_config_paths function to ensure all model checkpoint paths specified in the configuration files exist, preventing runtime errors. Added validation for the main model_path argument as well.
Configuration and Example Updates: Added new configuration files (ltx2_tp.json, ltx2_ulysses.json) and example shell scripts to demonstrate running LTX2 with Tensor Parallel and Ulysses (Sequence Parallel) settings. Various model checkpoint paths in existing configs and examples were updated for consistency.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces significant new features, including sequence parallelism, tensor parallelism, and VAE tiling for LTX2. The implementation of parallelism is extensive and touches many parts of the codebase. Overall, the changes are well-structured, particularly the use of wrapper classes for tensor parallelism. However, I've identified a critical bug in the weight loading logic for tensor parallelism that could lead to incorrect model loading. Additionally, there are several instances of hardcoded local paths and debug print statements that should be removed before merging. I also noted a potential regression in the image conditioning API, which has become less flexible. Please review my comments for detailed feedback and suggestions.

gemini-code-assist · 2026-01-23T06:21:57Z

lightx2v/models/networks/ltx2/model.py

+                else:
+                    # Non-TP weights or bias (bias is handled above if it's a TP weight's bias)
+                    # Skip bias keys that are already processed above
+                    if not (key.endswith(".bias") and key.replace(".bias", ".weight") in processed_weight_dict):


There's a critical bug in this condition. The check key.replace(".bias", ".weight") in processed_weight_dict will likely always be false for TP weights, because processed_weight_dict contains ranked keys (e.g., ...__tp_rank_0) for TP weights, not the original weight key. This can lead to biases of column-split TP weights being processed twice, causing incorrect model loading and behavior in tensor parallel mode.

A more robust, order-independent check is needed to see if the bias belongs to a column-split TP weight that has already been handled. I suggest checking the weight's properties directly.

Suggested change

if not (key.endswith(".bias") and key.replace(".bias", ".weight") in processed_weight_dict):

if not (key.endswith(".bias") and self._is_tp_weight(key.replace(".bias", ".weight")) and self._get_split_type(key.replace(".bias", ".weight")) == "col"):

lightx2v/models/runners/ltx2/ltx2_runner.py

gemini-code-assist · 2026-01-23T06:21:57Z

examples/ltx2/ltxt_i2av_distilled_fp8.py

+    model_path="/data/nvme0/gushiqiao/models/official_models/LTX-2/",
    model_cls="ltx2",
    task="i2av",
 )

 pipe.enable_quantize(
    dit_quantized=True,
-    dit_quantized_ckpt="Lightricks/LTX-2/ltx-2-19b-distilled-fp8.safetensors",
+    dit_quantized_ckpt="/data/nvme0/gushiqiao/models/official_models/LTX-2/ltx-2-19b-distilled-fp8.safetensors",


Hardcoded local paths are being used for model_path and dit_quantized_ckpt. This is not portable and will cause issues for other developers. It's recommended to use relative paths, environment variables, or a mechanism to load these from a local configuration file that is not committed to the repository.

gemini-code-assist · 2026-01-23T06:21:57Z

examples/ltx2/ltxt_t2av_distilled_fp8.py

+pipe = LightX2VPipeline(model_path="/data/nvme0/gushiqiao/models/official_models/LTX-2", model_cls="ltx2", task="t2av")

 pipe.enable_quantize(
    dit_quantized=True,
-    dit_quantized_ckpt="Lightricks/LTX-2/ltx-2-19b-distilled-fp8.safetensors",
+    dit_quantized_ckpt="/data/nvme0/gushiqiao/models/official_models/LTX-2/ltx-2-19b-distilled-fp8.safetensors",


This file contains hardcoded local paths for model_path and dit_quantized_ckpt. This makes the example not runnable for other users. Please consider using relative paths from the repository root, or loading paths from environment variables to improve portability.

lightx2v/common/ops/mm/mm_weight.py

lightx2v/models/networks/wan/model.py

gemini-code-assist · 2026-01-23T06:21:57Z

lightx2v/pipeline.py

-        print(config)
+        validate_config_paths(config)
        self.runner = self._init_runner(config)
+        print(self.runner.config)


This print statement seems to be for debugging. It should be removed before merging to keep the console output clean.

gemini-code-assist · 2026-01-23T06:21:57Z

scripts/wan/run_wan_i2v.sh

+lightx2v_path=/data/nvme0/gushiqiao/models/code/LightX2V
+model_path=/data/nvme0/gushiqiao/models/official_models/Wan2.1-I2V-14B-720P


The script has been updated with hardcoded local paths for lightx2v_path and model_path. This is not portable and should be reverted or replaced with a more generic solution, such as using placeholders or environment variables, to make the script usable by other developers.

[Feat] Support LTX2 sequence parallel, tensor parallel, and VAE decod…

aa243b2

…e tiling

fix

d52ce31

gemini-code-assist bot reviewed Jan 23, 2026

View reviewed changes

gushiqiao added 2 commits January 23, 2026 06:25

fix

cb3ae89

fix

dd2b5fd

llmc-reviewer approved these changes Jan 23, 2026

View reviewed changes

gushiqiao merged commit 746d44b into main Jan 23, 2026
2 checks passed

gushiqiao deleted the gsq/ltx-parallel branch January 23, 2026 06:29

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Feat] Support LTX2 sequence parallel, tensor parallel, and VAE decod… #810

[Feat] Support LTX2 sequence parallel, tensor parallel, and VAE decod… #810

Uh oh!

gushiqiao commented Jan 23, 2026

Uh oh!

gemini-code-assist bot commented Jan 23, 2026

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Jan 23, 2026

Uh oh!

Uh oh!

gemini-code-assist bot Jan 23, 2026

Uh oh!

gemini-code-assist bot Jan 23, 2026

Uh oh!

Uh oh!

Uh oh!

gemini-code-assist bot Jan 23, 2026

Uh oh!

gemini-code-assist bot Jan 23, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

	if not (key.endswith(".bias") and key.replace(".bias", ".weight") in processed_weight_dict):
	if not (key.endswith(".bias") and self._is_tp_weight(key.replace(".bias", ".weight")) and self._get_split_type(key.replace(".bias", ".weight")) == "col"):

		lightx2v_path=/data/nvme0/gushiqiao/models/code/LightX2V
		model_path=/data/nvme0/gushiqiao/models/official_models/Wan2.1-I2V-14B-720P

[Feat] Support LTX2 sequence parallel, tensor parallel, and VAE decod… #810

[Feat] Support LTX2 sequence parallel, tensor parallel, and VAE decod… #810

Uh oh!

Conversation

gushiqiao commented Jan 23, 2026

Uh oh!

gemini-code-assist bot commented Jan 23, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jan 23, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

gemini-code-assist bot Jan 23, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jan 23, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

gemini-code-assist bot Jan 23, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jan 23, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants