Implement Model Parallelism (Pipeline) for batch_size=1 & Fix Multi-GPU Device Mismatches

Browse Source

Feature (Pipeline Parallelism): implemented logic to handle batch_size=1 by splitting transformer blocks across GPUs.

Introduced ParallelBlock wrapper to manage execution across devices.

Data now flows in a pipeline sequence (e.g., GPU0 → GPU1 → ...) and automatically returns to the Lead Device for final processing.

Maintains original Data Parallel (batch splitting) behavior for batch_size > 1.

Fix (Device Mismatch): Resolved RuntimeError: Expected all tensors to be on the same device.

Updated ParallelBlock to recursively identify and move all input arguments (tensors, nested lists/tuples, dicts, and dataclasses) to the target device, rather than just the first argument.

Ensured auxiliary inputs (like timesteps and context) are correctly transferred alongside the main hidden states.

Fix (Syntax): Fixed a SyntaxError in clone_dataclass_or_object where a try block was missing its corresponding except.

Refactor: Improved tensor movement logic to be "deep" and recursive, preventing crashes when models use complex input structures.

...

This commit is contained in:

FearL0rd

2026-02-13 14:00:26 -06:00

committed by GitHub

parent 50b6fd1f74

commit 04907c683f

1 changed files with 184 additions and 184 deletions

Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL

Download Patch File Download Diff File Expand all files Collapse all files

any_device_parallel.py

+184 -184

File diff suppressed because it is too large Load Diff