instadeepai / instadeepai/Mava

[FEATURE] Remove Update Batch Size

Open
#1,132 0 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
935
Forks
123
Avg merge
6d 8h
Merged PRs (30d)
9

Description

### Please describe the purpose of the feature. Is it related to a problem?
This is somewhere between a feature and maintenance. Having to vmap an extra time is not necessarily as good an idea as we thought when we started rewriting Mava in JAX. I think it would be better to rather be able to use twice as many environments and then only have an outer `pmap`. It'll also remove a non-standard config option.

### Describe the solution you'd like
Remove the update_batch_size config option and update all systems accordingly.

### How do we know when implementation of this feature is complete?
Checklist:
- [ ] All systems are benchmarked and can get similar performance.

Contributor guide

Open the contributing guide

Research direction

Start by locating the update_batch_size configuration and its references across the systems in the Mava repository. Review how the current vmap and outer pmap flow is assembled, then benchmark every affected system. Done means the option is removed, all systems are updated, and benchmarks show similar performance.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning, performance
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.