Make sparse gradients configurable in IndexedMultiplier

[merajhashemi]

Summary

This PR adds a sparse_grad parameter to the IndexedMultiplier class to address compatibility issues with Distributed Data Parallel (DDP) training. The current implementation of IndexedMultiplier always uses sparse gradients, but sparse tensors are not supported in the all-reduce multi-GPU communication operation required by DDP (see pytorch/22400).

Changes

• Added a configurable sparse_grad parameter to IndexedMultiplier (default: True for backward compatibility).
• Updated documentation to explain the new parameter.

Usage

By making the sparsity configurable, users can now choose dense gradients (sparse_grad=False) when training with DDP.

from cooper.multipliers import IndexedMultiplier


# Current behavior (unchanged)
multiplier = IndexedMultiplier(num_constraints=100)

# New option for dense gradients
multiplier = IndexedMultiplier(num_constraints=100, sparse_grad=False)

Important Note

When using sparse_grad=False with stateful optimizers (e.g., Adam), optimizer states will be updated for all parameters. This may lead to incorrect optimization behavior since the optimizer will assume zero gradients for non-sampled indices when in reality these values should not be updated at all.

[github-actions]

Coverage report

Click to see where and how coverage changed

File	Statements	Missing	Coverage	Coverage (new stmts)	Lines missing
src/cooper/multipliers
multipliers.py
src/cooper/optim/torch_optimizers
nupi_optimizer.py					198
tests/multipliers
conftest.py
test_explicit_multipliers.py
tests/optim/torch_optimizers
test_nupi.py
Project Total

This report was generated by python-coverage-comment-action