Skip to content

Fix multi-node Ray GPU placement and launcher configuration - #2177

Open
rayg1234 wants to merge 3 commits into
mainfrom
rgao_placement_fix
Open

rayg1234 wants to merge 3 commits into
mainfrom
rgao_placement_fix

Conversation

@rayg1234

@rayg1234 rayg1234 commented Aug 24, 2026 •

Copy link
Copy Markdown
Contributor

Summary

  • propagate configurable workers-per-node and graph-parallel settings through the ASE dynamics prediction path
  • schedule multi-node prediction workers in one STRICT_SPREAD placement group and pin bundle 0 to the driver node, preventing duplicate physical GPU assignments
  • launch Ray from the active Python environment and forward Slurm GPUs-per-node and partition settings
  • This allow us to use multi-GPU ray on configurations that is not just 8 GPUs per node (ie: GB300)

Validation

  • pre-commit run --files on all six modified source files
  • Python bytecode compilation on all modified source files
  • successful 32-rank, 8-node all-to-all spatial graph-parallel dynamics canary

Targeted pytest tests were not run because pytest is not installed in the available fairchem virtual environment.

@meta-cla meta-cla Bot added the cla signed label Aug 24, 2026
@lbluque lbluque added patch Patch version release bug Something isn't working labels Aug 25, 2026
@rayg1234
rayg1234 requested review from lbluque and mshuaibii and removed request for lbluque August 26, 2026 00:22
lbluque
lbluque previously approved these changes Aug 26, 2026

@lbluque lbluque 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.

Minor comments and questions!

warmup_steps: int = 10,
workers: int = 0,
workers_per_node: int = 8,
gp_config: GraphParallelConfig | dict | None = None,

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.

allowing dicts is for backwards compat?

self.model_name = model_name
self.workers = workers
self.workers_per_node = workers_per_node
self.gp_config = gp_config

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.

Can we turn raw dicts into GPConfigs here to keep only that code path?

Comment on lines +471 to +473
sys.executable,
"-m",
"ray.scripts.scripts",

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.

Just curious why the switch directly to source rather than the cli entrypoint?

Comment on lines 771 to 773
bundle = {"CPU": workers}
if device == "cuda":
bundle["GPU"] = workers

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.

Question, why always 1:1 CPU to GPU ratio here?

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working cla signed patch Patch version release

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants