Lightning-AI / Lightning-AI/pytorch-lightning

Refactor the DeepSpeed strategy config management

Open
#17,472 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

fabric pl refactor strategy: deepspeed
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Outline & Motivation

DeepSpeed works by using a configuration file (dictionary) that allows customizing all of its aspects: https://www.deepspeed.ai/docs/config-json/

The `DeepSpeedStrategy` supports two ways of defining this:

1. Passing a config file, where every other argument becomes unused: https://github.com/Lightning-AI/lightning/blob/b792c90ea7148d61af192fde6c338ebbd355702f/src/lightning/fabric/strategies/deepspeed.py#L191
2. Exposes multiple of these arguments in the `__init__` that are used to define a base config. https://github.com/Lightning-AI/lightning/blob/b792c90ea7148d61af192fde6c338ebbd355702f/src/lightning/fabric/strategies/deepspeed.py#L242-L271

Option 2 is not scalable because:
- It forces us to duplicate all arguments
- Our docstrings might become outdated
- Our strategy defaults might diverge from the defaults in deepspeed
- It forces the user to either create an entire config or use these arguments
- Arguments might be different based on the installed deepspeed version as we support more than a single version.
- When deepspeed adds an argument that we don't expose, users have to switch to using the config

### Pitch

Remove all these exposed arguments and just have a `config` argument that overloads support for:
- Passing a path to a config file
```python
DeepSpeedStrategy(config="my/config/path.json")
```
- Passing a full config object
```python
config = ds.runtime.config.DeepSpeedConfig({"train_micro_batch_size_per_gpu": 2})
DeepSpeedStrategy(config=config)
```
- Passing a config dictionary (or a subset of it) that will update the default config
```python
config = {"zero_optimization": {"offload_optimizer": {"device": "cpu"}}}
DeepSpeedStrategy(config=config)
```

Where the default config is created by calling: https://github.com/microsoft/DeepSpeed/blob/085981bf1caf5d7d0b26d05f7c7e9487e1b35190/deepspeed/runtime/config.py#L674

### Additional context

DeepSpeed is considered experimental so we could do this breaking change: https://github.com/Lightning-AI/lightning/blob/b792c90ea7148d61af192fde6c338ebbd355702f/src/lightning/fabric/strategies/deepspeed.py#L99

cc @justusschock @awaelchli @carmocca

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in src/lightning/fabric/strategies/deepspeed.py, especially the DeepSpeedStrategy initializer and the existing config handling around the referenced lines. Review DeepSpeed's runtime/config.py defaults and determine how path, full config object, and partial dictionary inputs should be supported. Done means the exposed configuration arguments are removed and all three config forms update or use the default configuration as described.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
distributed-systems, machine-learning
Issue type
Refactor
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.