awslabs / awslabs/keys_values

Backward CPU offloading: Asynchronous transfer

Open
#62 16 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
18
Forks
8
PR merge metrics
No merged PRs in 30d

Description

**Is your feature request related to a problem? Please describe.**
CPU offloading in `backward` works by iterating this over layers:
- Load weights from CPU
- Run backward from head gradients
- Store gradients to CPU

These steps are currently run sequentially.

**Describe the solution you'd like**
Can we run these steps in parallel?

- Switch between two shards on GPU
- Run backward on one while loading weights for the other, and storing gradients for the previous

Needs clear understanding how async CPU <-> GPU transfer works! We know how transfer between GPUs works.

Contributor guide

Open the contributing guide

Research direction

Start by tracing the current backward CPU-offloading loop described in the issue: loading weights, running backward from head gradients, and storing gradients to CPU. Identify how the two GPU shards and CPU↔GPU transfers are currently handled; done means overlapping work across adjacent shards while preserving backward correctness.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.