[memories] Overlap shard transfers in async_serialize

[gspschmid]

Builds upon #22114, which introduces a fast-path for async checkpoint saving in which each array shard is copied through a single device-to-pinned-host transfer. So far, all of these transfers are serialized.

The present PR allows overlapping the transfers of a single array's shards. This is achieved by inserting asyncio.sleep(0) right after the transfer has been started, thereby permitting the _write_array(shard) coroutine to yield.

Why do we need the sleep? Python's await doesn't actually guarantee that control will be yielded to the eventloop so that other coroutines can be scheduled. Python ties coroutine semantics to that of generators. As I understand it, the only thing that causes control to pass to the eventloop (and thus allows other coroutines to be scheduled) is a yield. What we therefore need to achieve overlap in our situation, is something like

import types

@types.coroutine
def noop():
  yield

async def transfer_shard_to_host(shard):
  ...
  await noop() # behaves like asyncio.sleep(0)
  ...

This is, in fact, what asyncio.sleep(0) does.

An alternative, more roundabout way to achieve overlap would be to split _write_array(shard) into a first phase that initiates all the transfers, and provides the resulting host-resident arrays to a second phase, in which we invoke t.write(data) as before.

cc @yashk2810