Describe the bug
The data-parallel hybrid job example in 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
- From a Braket Notebook Instance, open
examples/hybrid_jobs/5_Parallelize_training_for_QML/Parallelize_training_for_QML.ipynb.
- Run cells sequentially up to and including the single-instance multi-GPU submission (
qml-dp1x-..., ml.g6.12xlarge, instance_count=1).
- After the cell returns, look up the job in the AWS Braket console (or via
aws braket get-job).
- 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:
Describe the bug
The data-parallel hybrid job example in Parallelize_training_for_QML.ipynb fails when executed against the current public Braket PyTorch jobs container. The jobs with
distribution="data_parallel"usewait_until_complete=Falseare submitted successfully, but the resulting jobs end in aFAILEDstate withAlgorithmError: ModuleNotFoundError: No module named 'smdistributed'. The training script (qml_script/train_dp.py) importssmdistributed.dataparallel.*, which is not present in the current PyTorch DLC on which the jobs container is based.To Reproduce
examples/hybrid_jobs/5_Parallelize_training_for_QML/Parallelize_training_for_QML.ipynb.qml-dp1x-...,ml.g6.12xlarge,instance_count=1).aws braket get-job).status: FAILEDwithfailureReasonreferencingModuleNotFoundError: No module named 'smdistributed'.Expected behavior
The hybrid job submitted with
distribution="data_parallel"runs to completion and reportsCOMPLETED.System information (please complete the following information as applicable):
amazon-braket-sdk==1.117.2,pennylane==0.45.0,pennylane-lightning[gpu]==0.45.0amazon-braket-pytorch-jobs:latestKnown workaround (if applicable)
None.
Additional context
The latest DLC images no longer include
smdistributed, so to continue supportingdata-paralleljobs 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: