Feature: Control number of vmapped envs in evaluator using arch.num_envs
#1071
Conversation
OmaymaMahjoub
requested review from
arnupretorius,
DriesSmit,
RuanJohn,
jcformanek,
siddarthsingh1,
sash-a,
ulricharmel,
callumtilbury and
WiemKhlifi
as code owners
There was a problem hiding this comment.
LGTM, thanks @OmaymaMahjoub for making things more flexible 🙏
| n_devices = len(jax.devices()) | ||
| episodes_per_device = config.arch.num_eval_episodes * eval_multiplier // n_devices |
There was a problem hiding this comment.
Sadly jax doesn't allow us to carry these fixed values (parallel_eval_batch_size) 😢
There was a problem hiding this comment.
I am quite sure there is a better way to implement this, but we can keep it as it is at the moment and create an issue calling for cleaning and readibility of the evaluator
mava/evaluator.py
Outdated
| parallel_eval_batch_size = min(config.arch.num_envs, episodes_per_device) | ||
| # Compute the number of sequential evaluation batches required per device | ||
| # to cover all episodes. | ||
| sequential_eval_batches = episodes_per_device // parallel_eval_batch_size |
There was a problem hiding this comment.
If you make keys of shape (sequential_eval_batches, num_vmapped_episodes) then you don't need the extra method, you can just call:
jax.lax.scan(jax.vmap(eval_one_episode), None, eval_states)I like this for two reasons, you don't have to calculate parallel_eval_batch_size twice and it is also clear that you're scanning and then vmapping over episodes. But this is a big change so if you don't have time it's ok. I think the evaluator is due for a big overhaul anyways
There was a problem hiding this comment.
I back this idea also 🔥
There was a problem hiding this comment.
I see so even the eval_init will be created inside the one episode, I will give it a try
Co-authored-by: Sasha <reallysasha@gmail.com>
What?
Modify the evaluator to limit the number of vmapped envs to
arch.num_envswhen the total number of evaluation episodesarch.num_eval_episodesexceeds this limit (instead of parallelise all thearch.num_eval_episodes). In such cases, evaluations are conducted in sequential batches, with each batch containingarch.num_envsparallel envs.Why?
Limiting parallel evaluations to num_envs prevents out-of-memory issues by avoiding vmap over all episodes at once.