[PR #29] [CLOSED] Adding c10d ops and additional aten ops for DTensor and FSDP2 #89

New issue

Closed

opened 2026-10-26 22:20:15 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 22:20:15 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/29
Author: @komikndr
Created: 2/6/2026
Status: ❌ Closed

Base: main ← Head: main

📝 Commits (10+)

e0321f1 c10d implementation for fp8
f6b8fb7 add additional ops require for sharding
bc5d7dd aten ops alias, and broadcast
d5538c6 Add scater ops, and rework copy to quantized dst to src format
1db1945 Add new_zeros op
da71067 pre and post gather fsdp_hook
c738713 Change README and NOTICE
bc2f646 Remove monkey patching for fsdp pre/post_all_gather
79ebc96 NVFP4 Enable
6bfe97b Update README for NVFP4

📊 Changes

5 files changed (+828 additions, -30 deletions)

View changed files

📝 NOTICE (+3 -0)
📝 README.md (+10 -0)
📝 comfy_kitchen/tensor/base.py (+59 -4)
📝 comfy_kitchen/tensor/fp8.py (+278 -5)
📝 comfy_kitchen/tensor/nvfp4.py (+478 -21)

📄 Description

Preliminary pull request adding FP8 operators to support Raylight custom nodes.

The c10d operators will not be touched during a normal Comfy run, sincen as far as I know only Raylight uses DTensor and its operators, which require an initialized distributed process group. Normal ComfyUI/main.py execution does not initialize this.

The operators being added:

c10d

all_gather_into_tensor
broadcast_
wait_tensor
scatter_

aten

slice
split
new_zeros
as_strided
cat
alias

In practice, for c10d only all_gather_into_tensor and wait_tensor are currently required, since wait_tensor is needed for the FSDP unshard path where an in-place copy cannot directly transfer a QuantizedTensor into an unquantized tensor.

To use this inside a model, the nn.Module cannot inherit from a "disable init weight" operator, since every sharded parameter must be properly accounted for.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/Comfy-Org/comfy-kitchen/pull/29 **Author:** [@komikndr](https://github.com/komikndr) **Created:** 2/6/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `main` --- ### 📝 Commits (10+) - [`e0321f1`](https://github.com/Comfy-Org/comfy-kitchen/commit/e0321f1bf53c69114556659983eae5f4e882e8a7) c10d implementation for fp8 - [`f6b8fb7`](https://github.com/Comfy-Org/comfy-kitchen/commit/f6b8fb70f536b390886f2bb4320da2aba569f1fb) add additional ops require for sharding - [`bc5d7dd`](https://github.com/Comfy-Org/comfy-kitchen/commit/bc5d7ddc6625c560dcb367c471de1f4597d06c16) aten ops alias, and broadcast - [`d5538c6`](https://github.com/Comfy-Org/comfy-kitchen/commit/d5538c62ab696480a940ce1b1407345379962fd6) Add scater ops, and rework copy to quantized dst to src format - [`1db1945`](https://github.com/Comfy-Org/comfy-kitchen/commit/1db1945296e08f583ef4a2d8f575640ee3d56773) Add new_zeros op - [`da71067`](https://github.com/Comfy-Org/comfy-kitchen/commit/da71067d17eae0b820fe545dae079b7f5623df40) pre and post gather fsdp_hook - [`c738713`](https://github.com/Comfy-Org/comfy-kitchen/commit/c738713425b247727dc23e667feec3100a80c69b) Change README and NOTICE - [`bc2f646`](https://github.com/Comfy-Org/comfy-kitchen/commit/bc2f6468e637d0d5816e06b54822e3f3e8f7121d) Remove monkey patching for fsdp pre/post_all_gather - [`79ebc96`](https://github.com/Comfy-Org/comfy-kitchen/commit/79ebc96301d5a1b1022271a0ee172becdf7633cd) NVFP4 Enable - [`6bfe97b`](https://github.com/Comfy-Org/comfy-kitchen/commit/6bfe97b160621455e727fa4783e7071fd9687a32) Update README for NVFP4 ### 📊 Changes **5 files changed** (+828 additions, -30 deletions) <details> <summary>View changed files</summary> 📝 `NOTICE` (+3 -0) 📝 `README.md` (+10 -0) 📝 `comfy_kitchen/tensor/base.py` (+59 -4) 📝 `comfy_kitchen/tensor/fp8.py` (+278 -5) 📝 `comfy_kitchen/tensor/nvfp4.py` (+478 -21) </details> ### 📄 Description Preliminary pull request adding FP8 operators to support Raylight custom nodes. The c10d operators will not be touched during a normal Comfy run, sincen as far as I know only Raylight uses DTensor and its operators, which require an initialized distributed process group. Normal `ComfyUI/main.py` execution does not initialize this. The operators being added: ### c10d - `all_gather_into_tensor` - `broadcast_` - `wait_tensor` - `scatter_` ### aten - `slice` - `split` - `new_zeros` - `as_strided` - `cat` - `alias` In practice, for c10d only `all_gather_into_tensor` and `wait_tensor` are currently required, since `wait_tensor` is needed for the FSDP unshard path where an in-place copy cannot directly transfer a `QuantizedTensor` into an unquantized tensor. To use this inside a model, the `nn.Module` cannot inherit from a "disable init weight" operator, since every sharded parameter must be properly accounted for. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-26 22:20:15 +00:00