amazon-braket / amazon-braket/amazon-braket-examples

Parallelize_training_for_QML hybrid job fails (ModuleNotFoundError: No module named 'smdistributed')

Open
#907 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
589
Forks
262
Avg merge
13h 38m
Merged PRs (30d)
7

Description

**Describe the bug**
The data-parallel hybrid job example in [Parallelize_training_for_QML.ipynb](https://github.com/amazon-braket/amazon-braket-examples/blob/main/examples/hybrid_jobs/5_Parallelize_training_for_QML/Parallelize_training_for_QML.ipynb) fails when executed against the current public Braket PyTorch jobs container. The jobs with `distribution="data_parallel"` use `wait_until_complete=False` are submitted successfully, but the resulting jobs end in a `FAILED` state with `AlgorithmError: ModuleNotFoundError: No module named 'smdistributed'`. The training script (`qml_script/train_dp.py`) imports `smdistributed.dataparallel.*`, which is not present in the current PyTorch DLC on which the jobs container is based.

**To Reproduce**
1. From a Braket Notebook Instance, open `examples/hybrid_jobs/5_Parallelize_training_for_QML/Parallelize_training_for_QML.ipynb`.
2. Run cells sequentially up to and including the single-instance multi-GPU submission (`qml-dp1x-...`, `ml.g6.12xlarge`, `instance_count=1`).
3. After the cell returns, look up the job in the AWS Braket console (or via `aws braket get-job`).
4. Observe `status: FAILED` with `failureReason` referencing `ModuleNotFoundError: No module named 'smdistributed'`.

**Expected behavior**
The hybrid job submitted with `distribution="data_parallel"` runs to completion and reports `COMPLETED`.

**System information (please complete the following information as applicable):**
- Package Versions: `amazon-braket-sdk==1.117.2`, `pennylane==0.45.0`, `pennylane-lightning[gpu]==0.45.0`
- Container: `amazon-braket-pytorch-jobs:latest`
- Working Environment: Braket Notebook Instance
- Region: us-east-1

**Known workaround (if applicable)**
None.

**Additional context**
The latest DLC images no longer include `smdistributed`, so to continue supporting `data-parallel` jobs like this we may need a different approach. This may require fixes in both this repo and the containers repo. Working branches here in both repos with a possible approach to fixing this:
- https://github.com/amazon-braket/amazon-braket-examples/compare/rmshaffer/data-parallel
- https://github.com/amazon-braket/amazon-braket-containers/compare/rmshaffer/data-parallel

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.