Skip to content

Adding support for diffusiongemma-26B-A4B-it - #456

Open
AnishPahilajani wants to merge 10 commits into
torch-spyre:mainfrom
AnishPahilajani:diffusion-gemma
Open

Adding support for diffusiongemma-26B-A4B-it#456
AnishPahilajani wants to merge 10 commits into
torch-spyre:mainfrom
AnishPahilajani:diffusion-gemma

Conversation

@AnishPahilajani

@AnishPahilajani AnishPahilajani commented Sep 1, 2026

Copy link
Copy Markdown
Contributor

You can run diffusion gemma on 2 spyre cards using the following command

torchrun --nproc-per-node 2 scripts/test_diffusion_gemma.py --model /models/diffusiongemma-26B-A4B-it --tp --max-new-tokens 256 --max-denoising-steps 8

output I get

 blue is due to a phenomenon called **Rayleigh scattering**.

Here is the step-by-step breakdown:

### 1. Sunlight is a rainbow
Although sunlight looks white to us, it is actually made up of all the colors of the rainbow (red, orange, yellow, green, blue, indigo, and violet). Each color travels as a wave with a different wavelength. Red light has long, lazy waves, while blue and violet light have short, choppy waves.

### 2. The atmosphere is an obstacle course
Earth’s atmosphere is filled with gases (mostly nitrogen and oxygen). These gas molecules. 

### 3. Scattering
When sunlight hits the atmosphere, it strikes the gas molecules and gets scattered in all directions. 
*   The **long waves** (red and yellow) pass through the atmosphere relatively undisturbed, hitting very molecules.
*   The **shorter waves** (blue and violet) strike the gas molecules and get scattered (scattered) in every direction.

### 4. Why do we see blue and not violet?
Since the blue light is being scattered in every part of the sky, when you look up, your eyes catch that scattered blue light coming from every part of the atmosphere

--- throughput ---
  max_denoising_steps : 8
  generated tokens    : 256
  wall time           : 266.74s
  throughput          : 1.0 tok/s

NOTE: throughput is dominated by Spyre↔CPU MoE round-trips.
  8 steps × 30 layers × 2 transfers/layer
  = 480 PCIe transfers per canvas.
  Reduce --max-denoising-steps to trade quality for speed.

@AnishPahilajani
AnishPahilajani force-pushed the diffusion-gemma branch 2 times, most recently from f561b8d to c54da7a Compare September 1, 2026 06:35
@AnishPahilajani AnishPahilajani changed the title adding support for diffusiongemma-26B-A4B-it Adding support for diffusiongemma-26B-A4B-it Sep 1, 2026
@AnishPahilajani
AnishPahilajani marked this pull request as ready for review September 1, 2026 17:39
@AnishPahilajani
AnishPahilajani force-pushed the diffusion-gemma branch 2 times, most recently from d0371a7 to 81ea9bd Compare September 2, 2026 02:54

@arielge arielge left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Haven't gone over the whole adapter logic, but generally - try to keep to a minimum changes to general API unless they are absolutely necessary

Comment thread hf_adapters/auto_spyre_model.py Outdated
Comment thread hf_adapters/auto_spyre_model.py Outdated
Comment thread hf_adapters/hf_diffusion_gemma.py Outdated
Comment thread hf_adapters/hf_common.py Outdated
@AnishPahilajani
AnishPahilajani force-pushed the diffusion-gemma branch 4 times, most recently from f45698b to 4c5ee71 Compare September 3, 2026 19:28
Comment thread hf_adapters/hf_diffusion_gemma.py Outdated
Comment thread hf_adapters/hf_diffusion_gemma.py Outdated
Comment thread hf_adapters/hf_diffusion_gemma.py Outdated
Comment thread tests/model_registry.py Outdated
"adapter": "hf_diffusion_gemma.py",
"size": "26b",
"dtype": "bfloat16",
"is_gated": True,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

model is not gated. Regardless, it should be explicitly excluded from the CAUSAL_PATHS, like we do for dspark entries

@AnishPahilajani
AnishPahilajani force-pushed the diffusion-gemma branch 2 times, most recently from 1da6c02 to 5678191 Compare September 9, 2026 17:32
@AnishPahilajani

Copy link
Copy Markdown
Contributor Author

fyi... 2 card does not work yet. So pls dont merge this.

torchrun --nproc_per_node=2 scripts/test_diffusion_gemma.py --model /models/diffusiongemma-26B-A4B-it --tp --dtype float16 --batch-size 2 --max-denoising-steps 8 --max-new-tokens 256

@BenjSz

BenjSz commented Sep 10, 2026

Copy link
Copy Markdown
Contributor

fyi... 2 card does not work yet. So pls dont merge this.

torchrun --nproc_per_node=2 scripts/test_diffusion_gemma.py --model /models/diffusiongemma-26B-A4B-it --tp --dtype float16 --batch-size 2 --max-denoising-steps 8 --max-new-tokens 256

So, for clarity, please move this PR to be draft.

@AnishPahilajani
AnishPahilajani marked this pull request as draft September 10, 2026 13:03
Signed-off-by: anishpahilajani <anish.pahilajani@ibm.com>
Signed-off-by: anishpahilajani <anish.pahilajani@ibm.com>
Signed-off-by: anishpahilajani <anish.pahilajani@ibm.com>
Signed-off-by: anishpahilajani <anish.pahilajani@ibm.com>
Signed-off-by: anishpahilajani <anish.pahilajani@ibm.com>
Signed-off-by: anishpahilajani <anish.pahilajani@ibm.com>
Signed-off-by: anishpahilajani <anish.pahilajani@ibm.com>
Signed-off-by: anishpahilajani <anish.pahilajani@ibm.com>
Signed-off-by: anishpahilajani <anish.pahilajani@ibm.com>
Signed-off-by: anishpahilajani <anish.pahilajani@ibm.com>
@AnishPahilajani
AnishPahilajani marked this pull request as ready for review September 10, 2026 18:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants