[platform]: add scripts and adapt for ascend, mlu, metal by Watebear · Pull Request #1208 · ModelTC/LightX2V

Watebear · 2026-07-02T05:38:14Z

No description provided.

# Conflicts: # configs/platforms/mlu/z_image_turbo_t2i.json # scripts/platforms/mlu/z_image_turbo_t2i.sh

gemini-code-assist

Code Review

This pull request introduces multi-platform support (Ascend NPU, Cambricon MLU, MetaX, and NVIDIA) for several text-to-image and text-to-video models, adding single-card and distributed (tensor-parallel and sequence-parallel) configurations along with platform-specific launcher scripts. Key changes include tensor-parallel weight loading for Flux2, sequence-parallel support for LongCat Image, and directory restructuring. The code review identified several critical issues: potential device mismatches when registering buffers on CPU, missing safety guards for uninitialized distributed environments in Flux2, a shape mismatch in HunyuanVideo's NPU attention output, an uninitialized sequence-parallel group attribute in LongCat Image, and invalid parameter usage in F.pad with replicate mode. Additionally, several test files and scripts contain broken file paths due to the directory restructuring and config renames, which need to be updated to prevent runtime and test failures.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

gemini-code-assist · 2026-07-02T05:41:42Z

+def _register_or_replace_buffer(module: torch.nn.Module, name: str, value: torch.Tensor) -> None:
+    if name in module._buffers:
+        module._buffers[name] = value
+    else:
+        module.register_buffer(name, value)


If _register_or_replace_buffer is called after the model has been moved to a non-CPU device (e.g., GPU or NPU), registering a CPU tensor as a buffer will keep it on the CPU. This can lead to device mismatch errors during the forward pass. We should dynamically cast the tensor to the device of the module's existing parameters or buffers.

Suggested change

def _register_or_replace_buffer(module: torch.nn.Module, name: str, value: torch.Tensor) -> None:

if name in module._buffers:

module._buffers[name] = value

else:

module.register_buffer(name, value)

def _register_or_replace_buffer(module: torch.nn.Module, name: str, value: torch.Tensor) -> None:

parameter = next(module.parameters(), None)

if parameter is not None:

value = value.to(parameter.device)

else:

buf = next(module.buffers(), None)

if buf is not None:

value = value.to(buf.device)

if name in module._buffers:

module._buffers[name] = value

else:

module.register_buffer(name, value)

gemini-code-assist · 2026-07-02T05:41:42Z

+    def _init_tensor_parallel(self):
+        if self.config.get("tensor_parallel", False):
+            self.use_tp = True
+            self.tp_group = self.config.get("device_mesh").get_group(mesh_dim="tensor_p")
+            self.tp_rank = dist.get_rank(self.tp_group)
+            self.tp_size = dist.get_world_size(self.tp_group)
+        else:
+            self.use_tp = False
+            self.tp_group = None
+            self.tp_rank = 0
+            self.tp_size = 1


If tensor_parallel is enabled but dist.is_initialized() is False or device_mesh is None, calling get_group or dist.get_rank will raise an error. We should add safety guards to ensure distributed is initialized and device_mesh is present before attempting to initialize tensor parallel.

Suggested change

def _init_tensor_parallel(self):

if self.config.get("tensor_parallel", False):

self.use_tp = True

self.tp_group = self.config.get("device_mesh").get_group(mesh_dim="tensor_p")

self.tp_rank = dist.get_rank(self.tp_group)

self.tp_size = dist.get_world_size(self.tp_group)

else:

self.use_tp = False

self.tp_group = None

self.tp_rank = 0

self.tp_size = 1

def _init_tensor_parallel(self):

if self.config.get("tensor_parallel", False) and dist.is_initialized() and self.config.get("device_mesh") is not None:

self.use_tp = True

self.tp_group = self.config.get("device_mesh").get_group(mesh_dim="tensor_p")

self.tp_rank = dist.get_rank(self.tp_group)

self.tp_size = dist.get_world_size(self.tp_group)

else:

self.use_tp = False

self.tp_group = None

self.tp_rank = 0

self.tp_size = 1

gemini-code-assist · 2026-07-02T05:41:42Z

+    out = torch.zeros(batch * seqlen, heads * dim, dtype=out_unpad.dtype, device=out_unpad.device)
+    out[indices] = out_unpad
+    return out.reshape(batch, seqlen, heads, dim)


The out_unpad tensor returned by npu_flash_attn has shape [total_unpadded_tokens, heads, dim]. Initializing out with shape [batch * seqlen, heads * dim] will cause a shape mismatch error when executing out[indices] = out_unpad. We should initialize out with shape [batch * seqlen, heads, dim] to match the unpadded tensor dimensions.

Suggested change

out = torch.zeros(batch * seqlen, heads * dim, dtype=out_unpad.dtype, device=out_unpad.device)

out[indices] = out_unpad

return out.reshape(batch, seqlen, heads, dim)

out = torch.zeros(batch * seqlen, heads, dim, dtype=out_unpad.dtype, device=out_unpad.device)

out[indices] = out_unpad

return out.reshape(batch, seqlen, heads, dim)

gemini-code-assist · 2026-07-02T05:41:42Z

        # Use transformer_in_channels to avoid conflict with VAE's in_channels
        self.in_channels = self.config.get("transformer_in_channels", self.config.get("in_channels", 64))
        self.attention_kwargs = {}
-        if self.config["seq_parallel"]:
-            raise NotImplementedError("Sequence parallel is not implemented for LongCatImageTransformerModel")
        self._init_infer_class()
        self._init_weights()
        self._init_infer()


When seq_parallel is enabled, _seq_parallel_pre_process and _seq_parallel_post_process access self.seq_p_group. However, self.seq_p_group is never initialized in LongCatImageTransformerModel.__init__, which will raise an AttributeError at runtime. We should initialize self.seq_p_group from the device mesh during model initialization.

# Use transformer_in_channels to avoid conflict with VAE's in_channels self.in_channels = self.config.get("transformer_in_channels", self.config.get("in_channels", 64)) self.attention_kwargs = {} if self.config.get("seq_parallel", False) and self.config.get("device_mesh") is not None: self.seq_p_group = self.config.get("device_mesh").get_group(mesh_dim="seq_p") else: self.seq_p_group = None self._init_infer_class() self._init_weights() self._init_infer()

gemini-code-assist · 2026-07-02T05:41:42Z

+def _pad(x: torch.Tensor, pad, mode: str = "constant", value: float = 0.0) -> torch.Tensor:
+    if mode == "replicate" and PLATFORM in _REPLICATE_PAD_FLOAT32_PLATFORMS and x.dtype == torch.bfloat16:
+        return F.pad(x.float(), pad, mode=mode, value=value).to(dtype=x.dtype)
+    return F.pad(x, pad, mode=mode, value=value)


In PyTorch, F.pad does not accept the value parameter when mode is "replicate". Passing value=value to F.pad with mode="replicate" will raise a ValueError. We should only pass value when mode is "constant".

Suggested change

def _pad(x: torch.Tensor, pad, mode: str = "constant", value: float = 0.0) -> torch.Tensor:

if mode == "replicate" and PLATFORM in _REPLICATE_PAD_FLOAT32_PLATFORMS and x.dtype == torch.bfloat16:

return F.pad(x.float(), pad, mode=mode, value=value).to(dtype=x.dtype)

return F.pad(x, pad, mode=mode, value=value)

def _pad(x: torch.Tensor, pad, mode: str = "constant", value: float = 0.0) -> torch.Tensor:

if mode == "replicate" and PLATFORM in _REPLICATE_PAD_FLOAT32_PLATFORMS and x.dtype == torch.bfloat16:

return F.pad(x.float(), pad, mode=mode).to(dtype=x.dtype)

if mode == "constant":

return F.pad(x, pad, mode=mode, value=value)

return F.pad(x, pad, mode=mode)

gemini-code-assist · 2026-07-02T05:41:42Z

+def test_hy15_dist8_config_and_launcher_default_to_single_card_alignment():
+    config = json.loads(Path("configs/platforms/nvidia/dist_8/hunyuan_video_t2v_480p.json").read_text(encoding="utf-8"))
+    script = Path("scripts/platforms/nvidia/dist_8/run_hy15_t2v_480p.sh").read_text(encoding="utf-8")
+
+    assert config["align_single_card_shape"] is True
+    assert "gpus=${GPUS:-8}" in script


The test file references configs/platforms/nvidia/dist_8/hunyuan_video_t2v_480p.json and scripts/platforms/nvidia/dist_8/run_hy15_t2v_480p.sh. However, the directory was renamed to dist and the config file was renamed to hunyuan_video_t2v_480p_dist8.json. This will cause a FileNotFoundError when running the tests. We should update the paths to match the new directory structure.

Suggested change

def test_hy15_dist8_config_and_launcher_default_to_single_card_alignment():

config = json.loads(Path("configs/platforms/nvidia/dist_8/hunyuan_video_t2v_480p.json").read_text(encoding="utf-8"))

script = Path("scripts/platforms/nvidia/dist_8/run_hy15_t2v_480p.sh").read_text(encoding="utf-8")

assert config["align_single_card_shape"] is True

assert "gpus=${GPUS:-8}" in script

def test_hy15_dist8_config_and_launcher_default_to_single_card_alignment():

config = json.loads(Path("configs/platforms/nvidia/dist/hunyuan_video_t2v_480p_dist8.json").read_text(encoding="utf-8"))

script = Path("scripts/platforms/nvidia/dist/run_hy15_t2v_480p.sh").read_text(encoding="utf-8")

assert config["align_single_card_shape"] is True

assert "gpus=${GPUS:-8}" in script

gemini-code-assist · 2026-07-02T05:41:42Z

+def test_z_image_nvidia_dist_uses_head_divisible_seq_parallel():
+    config = json.loads(Path("configs/platforms/nvidia/dist_8/z_image_turbo_t2i.json").read_text(encoding="utf-8"))
+    script = Path("scripts/platforms/nvidia/dist_8/z_image_turbo_t2i.sh").read_text(encoding="utf-8")
+
+    assert config["enable_cfg"] is False
+    assert config["parallel"] == {
+        "seq_p_size": 2,
+        "seq_p_attn_type": "ulysses",
+    }
+    assert "gpus=${GPUS:-2}" in script
+    assert 'CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-4,5}"' in script


The test file references configs/platforms/nvidia/dist_8/z_image_turbo_t2i.json and scripts/platforms/nvidia/dist_8/z_image_turbo_t2i.sh. However, the directory was renamed to dist and the config file was renamed to z_image_turbo_t2i_dist2.json. This will cause a FileNotFoundError when running the tests. We should update the paths to match the new directory structure.

Suggested change

def test_z_image_nvidia_dist_uses_head_divisible_seq_parallel():

config = json.loads(Path("configs/platforms/nvidia/dist_8/z_image_turbo_t2i.json").read_text(encoding="utf-8"))

script = Path("scripts/platforms/nvidia/dist_8/z_image_turbo_t2i.sh").read_text(encoding="utf-8")

assert config["enable_cfg"] is False

assert config["parallel"] == {

"seq_p_size": 2,

"seq_p_attn_type": "ulysses",

}

assert "gpus=${GPUS:-2}" in script

assert 'CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-4,5}"' in script

def test_z_image_nvidia_dist_uses_head_divisible_seq_parallel():

config = json.loads(Path("configs/platforms/nvidia/dist/z_image_turbo_t2i_dist2.json").read_text(encoding="utf-8"))

script = Path("scripts/platforms/nvidia/dist/z_image_turbo_t2i.sh").read_text(encoding="utf-8")

assert config["enable_cfg"] is False

assert config["parallel"] == {

"seq_p_size": 2,

"seq_p_attn_type": "ulysses",

}

assert "gpus=${GPUS:-2}" in script

assert 'CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-4,5}"' in script

gemini-code-assist · 2026-07-02T05:41:42Z

+def test_longcat_nvidia_dist8_launcher_uses_torchrun():
+    script = Path("scripts/platforms/nvidia/dist_8/longcat_image_t2i.sh").read_text(encoding="utf-8")
+
+    assert 'torchrun --nproc_per_node="${gpus}" -m lightx2v.infer' in script
+    assert "gpus=${GPUS:-8}" in script
+    assert "python -m lightx2v.infer" not in script
+
+
+def test_longcat_nvidia_dist8_config_matches_qwen_image_parallel_layout():
+    with open("configs/platforms/nvidia/dist_8/longcat_image_t2i.json", "r", encoding="utf-8") as f:
+        config = json.load(f)
+
+    assert config["parallel"] == {
+        "seq_p_size": 4,
+        "seq_p_attn_type": "ulysses",
+        "cfg_p_size": 2,
+    }


The test file references configs/platforms/nvidia/dist_8/longcat_image_t2i.json and scripts/platforms/nvidia/dist_8/longcat_image_t2i.sh. However, the directory was renamed to dist and the config file was renamed to longcat_image_t2i_dist8.json. This will cause a FileNotFoundError when running the tests. We should update the paths to match the new directory structure.

Suggested change

def test_longcat_nvidia_dist8_launcher_uses_torchrun():

script = Path("scripts/platforms/nvidia/dist_8/longcat_image_t2i.sh").read_text(encoding="utf-8")

assert 'torchrun --nproc_per_node="${gpus}" -m lightx2v.infer' in script

assert "gpus=${GPUS:-8}" in script

assert "python -m lightx2v.infer" not in script

def test_longcat_nvidia_dist8_config_matches_qwen_image_parallel_layout():

with open("configs/platforms/nvidia/dist_8/longcat_image_t2i.json", "r", encoding="utf-8") as f:

config = json.load(f)

assert config["parallel"] == {

"seq_p_size": 4,

"seq_p_attn_type": "ulysses",

"cfg_p_size": 2,

}

def test_longcat_nvidia_dist8_launcher_uses_torchrun():

script = Path("scripts/platforms/nvidia/dist/longcat_image_t2i.sh").read_text(encoding="utf-8")

assert 'torchrun --nproc_per_node="${gpus}" -m lightx2v.infer' in script

assert "gpus=${GPUS:-8}" in script

assert "python -m lightx2v.infer" not in script

def test_longcat_nvidia_dist8_config_matches_qwen_image_parallel_layout():

with open("configs/platforms/nvidia/dist/longcat_image_t2i_dist8.json", "r", encoding="utf-8") as f:

config = json.load(f)

assert config["parallel"] == {

"seq_p_size": 4,

"seq_p_attn_type": "ulysses",

"cfg_p_size": 2,

}

gemini-code-assist · 2026-07-02T05:41:42Z

+    --model_path "${model_path}" \
+    --config_json "${lightx2v_path}/configs/platforms/nvidia/dist_8/hunyuan_video_t2v_480p.json" \
+    --prompt "A close-up shot captures a scene on a polished, light-colored granite kitchen counter, illuminated by soft natural light from an unseen window. Initially, the frame focuses on a tall, clear glass filled with golden, translucent apple juice standing next to a single, shiny red apple with a green leaf still attached to its stem. The camera moves horizontally to the right. As the shot progresses, a white ceramic plate smoothly enters the frame, revealing a fresh arrangement of about seven or eight more apples, a mix of vibrant reds and greens, piled neatly upon it. A shallow depth of field keeps the focus sharply on the fruit and glass, while the kitchen backsplash in the background remains softly blurred. The scene is in a realistic style." \


The script references configs/platforms/nvidia/dist_8/hunyuan_video_t2v_480p.json, which does not exist because the directory was renamed to dist and the file was renamed to hunyuan_video_t2v_480p_dist8.json. This will cause a FileNotFoundError when executing the script. We should update the path to the correct config file.

Suggested change

--model_path "${model_path}" \

--config_json "${lightx2v_path}/configs/platforms/nvidia/dist_8/hunyuan_video_t2v_480p.json" \

--prompt "A close-up shot captures a scene on a polished, light-colored granite kitchen counter, illuminated by soft natural light from an unseen window. Initially, the frame focuses on a tall, clear glass filled with golden, translucent apple juice standing next to a single, shiny red apple with a green leaf still attached to its stem. The camera moves horizontally to the right. As the shot progresses, a white ceramic plate smoothly enters the frame, revealing a fresh arrangement of about seven or eight more apples, a mix of vibrant reds and greens, piled neatly upon it. A shallow depth of field keeps the focus sharply on the fruit and glass, while the kitchen backsplash in the background remains softly blurred. The scene is in a realistic style." \

--model_path "${model_path}" \

--config_json "${lightx2v_path}/configs/platforms/nvidia/dist/hunyuan_video_t2v_480p_dist8.json" \

--prompt "A close-up shot captures a scene on a polished, light-colored granite kitchen counter, illuminated by soft natural light from an unseen window. Initially, the frame focuses on a tall, clear glass filled with golden, translucent apple juice standing next to a single, shiny red apple with a green leaf still attached to its stem. The camera moves horizontally to the right. As the shot progresses, a white ceramic plate smoothly enters the frame, revealing a fresh arrangement of about seven or eight more apples, a mix of vibrant reds and greens, piled neatly upon it. A shallow depth of field keeps the focus sharply on the fruit and glass, while the kitchen backsplash in the background remains softly blurred. The scene is in a realistic style." \

gemini-code-assist · 2026-07-02T05:41:42Z

+    --task t2i \
+    --model_path "${model_path}" \
+    --config_json "${lightx2v_path}/configs/platforms/nvidia/dist/z_image_turbo_t2_dist2.json" \


The script references configs/platforms/nvidia/dist/z_image_turbo_t2_dist2.json, but the added config file is named z_image_turbo_t2i_dist2.json (with an i after t2). This will cause a FileNotFoundError when executing the script. We should update the path to the correct config file.

Suggested change

--task t2i \

--model_path "${model_path}" \

--config_json "${lightx2v_path}/configs/platforms/nvidia/dist/z_image_turbo_t2_dist2.json" \

--model_path "${model_path}" \

--config_json "${lightx2v_path}/configs/platforms/nvidia/dist/z_image_turbo_t2i_dist2.json" \

--prompt 'Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. Neon lightning-bolt lamp (⚡️), bright yellow glow, above extended left palm. Soft-lit outdoor night background, silhouetted tiered pagoda (西安大雁塔), blurred colorful distant lights.' \

Watebear and others added 16 commits June 19, 2026 12:26

[ascend]: add scripts and adapt codes for ascend npu

52e3d0f

[metax]: add scripts and adapt codes for metax

f5aca0b

[mlu]: add scripts and adapt codes in mlu

f1aa944

[scripts]: nv compare

1db8b73

Merge remote-tracking branch 'origin/main' into ascend

841009f

# Conflicts: # configs/platforms/mlu/z_image_turbo_t2i.json # scripts/platforms/mlu/z_image_turbo_t2i.sh

支持flux2-dev的tp

3619056

[feat]: flux2-dev support tp

f63b568

[feat]: loncat-image support seq_parallel

1a3f984

[nvidia]: scripts

2788644

Add NVIDIA distributed platform scripts

1674697

Add NVIDIA distributed platform configs

3afb621

Update Wan MoE NVIDIA benchmark resolution

313f82f

Update Ascend NPU platform configs and scripts

bdd5c08

Reorganize MLU platform configs and scripts

22809e8

update metax script

4d6680c

rm tests

1e505a1

gemini-code-assist Bot reviewed Jul 2, 2026

View reviewed changes

rm

289f46d

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[platform]: add scripts and adapt for ascend, mlu, metal#1208

[platform]: add scripts and adapt for ascend, mlu, metal#1208
Watebear wants to merge 17 commits into
mainfrom
ascend

Watebear commented Jul 2, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

Watebear commented Jul 2, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant