Skip to content

Fix zero-worker prediction on MPS-capable machines - #279

Open
leeshengis wants to merge 1 commit into
Unbabel:masterfrom
leeshengis:fix-mps-zero-worker-predict
Open

leeshengis wants to merge 1 commit into
Unbabel:masterfrom
leeshengis:fix-mps-zero-worker-predict

Conversation

@leeshengis

Copy link
Copy Markdown

Problem

COMET prediction with gpus=0 defaults to num_workers=0. On MPS-capable machines, predict still passes multiprocessing_context="fork" to PyTorch DataLoader. DataLoader rejects that combination before inference starts

This matches the existing M1 report and M4 confirmation in huggingface/evaluate#673. Evaluate forwards its CPU configuration to COMET, so this fixes the underlying COMET path. This PR does not claim to close the issue in the separate evaluate repository

Change

Only supply the fork context when MPS is available and the resolved worker count is positive. Preserve the current positive-worker MPS behavior and the non-MPS default

Add a unittest exercising the real CometModel.predict and real DataLoader construction, without downloading a model

Validation

  • Before the change: two of eight subcases raise the reported DataLoader ValueError
  • After the change: all eight subcases pass
  • Positive workers retain fork on MPS; non-MPS cases retain None
  • Test file passes black (80 columns, py38 target) and isort
  • Independent review found no correctness blockers and reran the regression successfully

Command:

HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 python -m unittest discover -s tests/unit -p test_predict_dataloader_context.py -v

Tested on Linux x86_64, Python 3.12.14, CPU torch 2.14.1+cpu, COMET 2.2.7. The tested baseline source matches master da61b81

MPS availability is simulated; the model instance and Lightning Trainer are mocked. The real DataLoader constructor validates the parameters, but this is not a native Mac, GPU, worker-iteration or model-inference test. The full model-dependent suite was not run. The test skips Windows. Current upstream CI runs a CLI smoke check, not unit tests

AI assistance was used to prepare the patch and tests and to review and run this focused regression. No claim of human review or native Apple Silicon validation is made

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

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant