deepspeedai / deepspeedai/DeepSpeedExamples
Deepspeed-chat feature request: specify which epoch/iteration to start EMA
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 6.8k
- Forks
- 1.1k
- Avg merge
- 2d 16h
- Merged PRs (30d)
- 1
Description
leveraging learnings from the mosaic ml blogpost on speeding up diffusion model training https://www.mosaicml.com/blog/diffusion it doesn't seem like we need to update the EMA model weights at every iteration of training, the last <5% of steps should be fine.
i propose adding an argument like --ema_start_pct with a default like 0.95, so start_ema_iter = int(args.ema_start_pct *num_total_iters), then keeping a global track of current_iter` such that
if args.enable_ema and current_iter >= start_ema_iter:
moving_average(rlhf_engine.actor,
rlhf_engine.actor_ema,
zero_stage=args.actor_zero_stage)
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in the DeepSpeed-chat training entry point by locating argument parsing, the total-iteration calculation, current_iter tracking, and the moving_average call on rlhf_engine.actor and actor_ema. Add the requested EMA start percentage behavior and verify that EMA updates begin only at the configured iteration, including the default value.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100