[GH-ISSUE #28] Is it okay to quantized dst based on src layout in handle_copy base.py ? #15

New issue

Closed

opened 2026-10-26 22:19:08 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 22:19:08 +00:00

Owner

Originally created by @komikndr on GitHub (Feb 2, 2026).
Original GitHub issue: https://github.com/Comfy-Org/comfy-kitchen/issues/28

Preamble: I want to add multiple ops to enable FSDP2 + the DCP API for Raylight (parallel stuff).

Commit:
github.com/komikndr/comfy-kitchen-distributed@d5538c62ab

Many new ops have been written in their own Layout classes, and I think they should not touch the core ComfyUI forward path, since most of them are c10d operators (e.g., broadcast, all_gather, wait_tensor). These will not be invoked unless someone uses the DTensor API, which as far as I know, is currently only used by Raylight.

However, there is a rewrite in handle_copy in base.py. Although it already checks whether the QuantizedTensor shares the same layout, I added a new case: if dst is not a QuantizedTensor (i.e., it is a native PyTorch tensor), it will be quantized to match the source type. This change is required since FSDP wait_for_unshard need to alloc some output buffer implied by QuantizedTensor.dtype

Will this break anything down the line? My preliminary tests have not caught any errors yet.

Thanks for your time 👍

Originally created by @komikndr on GitHub (Feb 2, 2026). Original GitHub issue: https://github.com/Comfy-Org/comfy-kitchen/issues/28 **Preamble:** I want to add multiple ops to enable FSDP2 + the DCP API for Raylight (parallel stuff). Commit: https://github.com/komikndr/comfy-kitchen-distributed/commit/d5538c62ab696480a940ce1b1407345379962fd6 Many new ops have been written in their own Layout classes, and I think they should not touch the core ComfyUI forward path, since most of them are c10d operators (e.g., `broadcast`, `all_gather`, `wait_tensor`). These will not be invoked unless someone uses the DTensor API, which as far as I know, is currently only used by Raylight. However, there is a rewrite in `handle_copy` in `base.py`. Although it already checks whether the QuantizedTensor shares the same layout, I added a new case: if `dst` is not a QuantizedTensor (i.e., it is a native PyTorch tensor), it will be quantized to match the source type. This change is required since FSDP `wait_for_unshard` need to alloc some output buffer implied by QuantizedTensor.dtype Will this break anything down the line? My preliminary tests have not caught any errors yet. Thanks for your time 👍

bsakaguchi closed this issue

2026-10-26 22:19:08 +00:00

bsakaguchi referenced this issue

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

[PR #15] [MERGED] Add Development.SABIModule to Python package search #81