feat: switch sebublba to using shard_map like mava

[sash-a]

What?

Some mava upgrades, as promised 😄
Quite a few nice changes by upgrading to shard_map over pmap here, which avoids some unnecessary device transfers. I found these in mava using transfer guard, quite a nice tool.

Why?

To make sebulba go brrr

How?

• Switch to shard_map
• Remove the default actor device - this causing some unnecessary transfers in the pipeline
• Stop using flax replicate/unreplicate rather explicitly put
• Move a block_until_ready from the params source to the learner. I think the unreplicate that in the learner before was doing this, without this block we get weird and undefined seg faults
• One issue I see is that we're now passing the same key to all the learners, we are actually doing this in mava also and I realize it is a minor issue, I'm not entirely sure how to fix it, I tried quickly to switch the sharding for the key to the data sharding which I think should fix it, but it didn't...it's Sunday, hopefully I have time to look at it during the week or if you could find a solution that would be awesome

NOTE:

This is very much not benchmarked. I pulled in Mava's changes in about an hour and I tested it locally and it solves cartpole, but I haven't checked on a TPU or with a harder env

[EdanToledo]

Thanks so much, I'll try review this and test it tomorrow on a GPU.

[EdanToledo]

This is the only thing I am not sure of. Why do we put the initial params on the first actor device instead of the cpu or something.

[EdanToledo]

I just did a comparison, and it seems like sebulba on main is faster. Looking at all the timing statistics, its the pipeline that is slowing things down. Everything else on this branch is faster. I am trying this on a 2-gpu system. I'll try figure out why its slow.

[EdanToledo]

I just did a comparison, and it seems like sebulba on main is faster. Looking at all the timing statistics, its the pipeline that is slowing things down. Everything else on this branch is faster. I am trying this on a 2-gpu system. I'll try figure out why its slow.

So basically, because the actors are faster now, the pipeline fills up faster and then causes the insertions to slow down a lot i think because there is more waiting and blocking which i imagine increases the overheads. I'm not sure on the solution to it though.

[sash-a]

I think CPU for this is best because it should never block actors or learners, but not sure because it would be much faster on an accelerator 🤔

[sash-a]

So basically, because the actors are faster now, the pipeline fills up faster and then causes the insertions to slow down a lot i think because there is more waiting and blocking which i imagine increases the overheads. I'm not sure on the solution to it though.

I think faster actors should almost always be better, maybe it just needs some hyper parameter tuning?