JetBrains-Research / opacus / 19234922926

78%

jbr-main: 81%

push

github

Evgri243/multi device models (#796)

Summary:
## Types of changes

- [x] Bug fix (non-breaking change which fixes an issue)
- [x] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to change)
- [ ] Docs change / refactoring / dependency upgrade

## Motivation and Context / Related issue

This PR adds support for multi-device training scenarios where model parameters are distributed across multiple GPU devices (e.g., when assigning different layers directly with `module.to(device[I])` oe using `device_map="auto"` with accelerate).

**Problem solved:**
When training large models that don't fit on a single GPU, parameters and gradients can be spread across multiple devices. The existing Opacus optimizers and gradient clipping modules assumed all tensors were on the same device, causing runtime errors during norm computation and gradient clipping operations.

**Changes:**
1. **Sequential multi-device execution support (https://github.com/meta-pytorch/opacus/issues/9)**: Modified `DPOptimizer` and `AdaClipDPOptimizer` to move tensors to appropriate devices before operations like `torch.stack()` and `torch.einsum()`, preventing device mismatch errors during gradient clipping and accumulation.

2. **Multi-device support in GradSampleModuleFastGradientClipping (https://github.com/meta-pytorch/opacus/issues/10)**: Extended multi-device handling to `GradSampleModuleFastGradientClipping`, `DPPerLayerOptimizer`, and additional edge cases in optimizers that were previously uncovered.

## How Has This Been Tested

- The code was used to train 7B Zetta model with LoRA on 8xH200 GPU node.
- Added test suite in `multidevice_optimizer_test.py` covering:
  - `DPOptimizer`, `AdaClipDPOptimizer`, and `DPPerLayerOptimizer` with multi-device models
  - Both `clip_and_accumulate()` and full `step()` operations
  - Helper function `_clip_and_accumulate_parameter()` with multi-device paramete... (continued)

40 of 252 new or added lines in 5 files covered. (15.87%)

5672 of 7256 relevant lines covered (78.17%)

1.74 hits per line