Lightning-AI / Lightning-AI/pytorch-lightning

Add additional parameters to the DeepSpeedStrategy

Open
#19,278 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

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

Description

### Description & Motivation

Related: #12820

Actually, the description is pretty much the same of that issue:
> .. a few more parameters have become available for DeepSpeed including ignore_unused_parameters (the opposite of find_unused_parameters for DDP).
>
> https://www.deepspeed.ai/docs/config-json/#zero-optimizations-for-fp16-training
>
> Integrate ignore_unused_parameters, round_robin_gradients and stage3_gather_16bit_weights_on_model_save into the Strategy.

### Pitch

An ultimate solution might be to allow `.json`-based configuration as done in Deepspeed, but adding more parameters would be a great remedy at least in the short term. Personally I'm interested in `ignore_unused_parameters`, which might become more and more useful as people train multimodal LLMs (related: https://github.com/huggingface/accelerate/issues/2194)

### Alternatives

Implementing this on my Lightning fork and use it.

### Additional context

_No response_

cc @borda @awaelchli

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 by locating the DeepSpeedStrategy entry point and read related issue #12820 alongside the DeepSpeed configuration documentation linked in the issue. Done means the Strategy supports ignore_unused_parameters, round_robin_gradients, and stage3_gather_16bit_weights_on_model_save, with tests covering the added parameters.

Written by the indexing model from the issue text.

Assessment

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.