Skip to content

fix: --device autodetect the available backend - #418

Open
rizzler13 wants to merge 1 commit into
kyutai-labs:mainfrom
rizzler13:fix/auto-device-detection
Open

fix: --device autodetect the available backend#418
rizzler13 wants to merge 1 commit into
kyutai-labs:mainfrom
rizzler13:fix/auto-device-detection

Conversation

@rizzler13

Copy link
Copy Markdown

Closes #396

Changes added:
Adds a small _get_default_device() helper to each of the 3 CLI entry points that auto-detects the best available backend:

  • torch.cuda.is_available()"cuda"
  • torch.backends.mps.is_available()"mps" (Apple Silicon via Metal)
  • fallback → "cpu"

Files changed:

  • moshi/moshi/server.py — added _get_default_device(), updated --device default, guarded torch.cuda.synchronize()
  • moshi/moshi/run_inference.py — added _get_default_device(), updated --device default
  • moshi/moshi/run_tts.py — added _get_default_device(), updated --device default

Explicit --device cuda still works exactly as before. The only behavioral change is the default when --device is omitted.

Note: seed_all() is duplicated identically in server.py and run_inference.py. It's already correctly guarded (if torch.cuda.is_available()) so it's not part of the crash path, I left it untouched to keep this PR focused. Happy to consolidate it in a follow up if that'd be useful.

@nouraellm nouraellm left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

LGTM! The only thing I'd verify is that the existing default dtype (torch.bfloat16) behaves correctly when the auto-selected device is mps or cpu. Otherwise I don't see any issues with the change.

@rizzler13

Copy link
Copy Markdown
Author

Just double checked this on my machine to be sure,
both cpu and mps work fine with the default configs (bfloat16/float16), and nn.Linear runs without any issues.

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.

Add automatic device detection to avoid CUDA assumption errors

2 participants