[tensor] ColoTensor supports ZeRo

[ver217]

Usage:

chunk_size = 38 * 1024**2 if use_chunk else None
chunk_manager = ChunkManager(chunk_size, enable_distributed_storage=use_zero)
model = ColoDDPV2(model, chunk_manager)

chunk_size=None means chunk is not used.

[feifeibear]

self._chunk_manager
use _XXX as an internal var of Class.

[Wesley-Jzy]

typo: tensor

[ver217]

fixed

[feifeibear]

The main concern is about suspending parameters.
For example

class Net(_)
  def __init__
       self.fc1 = Linear()
       self.suspend_param = torch.Paramer(..)

The self.suspend_param will occur in module.parameters(). So your DDPV2 will append it to the chunk manager.
Managing the state of the param will be a disaster for your design?

[ver217]

The main concern is about suspending parameters. For example

class Net(_)
  def __init__
       self.fc1 = Linear()
       self.suspend_param = torch.Paramer(..)

The self.suspend_param will occur in module.parameters(). So your DDPV2 will append it to the chunk manager. Managing the state of the param will be a disaster for your design?

Finally, we will use computation graph to derive the computation order of params and filter those unsed params. Chunk manager only manages used params

[feifeibear]

The name update is not informative enough.
I think you mean
def copy_tensor_to_chunk_slice()

[ver217]

done

[feifeibear]

is the access_chunks necessary?
It is only used for this line. You can know whether the chunk is accessed via its tensor states?

[ver217]

If a rank store a chunk, its initial state is HOLD, but the rank should broadcast.

[feifeibear]

Cannot understand what you mean? You can say Chinese here.

[ver217]

For example, dp size = 2 here, and rank0 stores chunk0. The initial state of chunk0 in rank0 is HOLD, which in rank1 is FREE. We want to access chunk0 now, even rank0 has chunk0, rank0 have to do broadcast(). We can determine whether a rank has a chunk by state, but we cannot know whether the broadcast() is done by state.

[feifeibear]

what does it mean?
reusing grad fp16 with param fp16?
should it be p.data = p.grad?

[ver217]

Set p.grad to correct ptr here. p.data saves grad, due to reuse.

[feifeibear]

if a chunk is moved from gpu to cpu later. This line makes p.grad point to an old memory space.

[ver217]

No, move device will update p.data at the same time.

[feifeibear]

This line makes p.grad -> p.data (addr1)
Afterwards, you move chunk
let p.data (addr2)
However, p.grad still points to addr1.

[ver217]

If so, just set p.grad again, or move this code snippets to optimizer.step() after device moving is done.

[ver217]

To check grad in unit test, I just set p.grad here.

[feifeibear]

Is it necessary to build a dict {param: chunk slice} to index grad and its true memory space (may reuse with param.data} ?

[ver217]

Not necessary now I think. p.grad should always point to the chunk slice memory of p, as we always reuse now. If not reuse, I think it's necessary.

[feifeibear]

OK, add it if necessary later.

[feifeibear]

OK, add it if necessary later.