Write only parameter groups that have changed rather than all parameter groups in the checkpoint

[nkandpa2]

Write only parameter groups that have changed rather than all parameter groups in the checkpoint

[blester125]

There was talk of using the metadata file to implement this to avoid reading both the checkpoint file and the git-theta model dir. This is a great idea. I have a suggestion on the API/implementation for this through.

Instead of something like looping through the model parameters and comparing them to the metadata file, we can have the git-theta model dir reading be lazy. Instead of returning actual parameter values it returns some object with tensor meta data (shape, dtype, hash, and maybe even a __eq__ function that knows how to hash/compare to a real tensor [I got myself excited with this lol]!) with methods to read and replace itself with the real tensor.

With this API, comparisions between the checkpoint tree and the model dir tree would look just like any tree intersection instead of being custom.

It could be overkill as this is similar to an approach we have in t5x where the main use of lazy arrays is for checkpoint manipulation for huge models which might not be a use case we care about now (but could be useful for parameter-efficient update applications in the future).

[blester125]

We talked in coworking and the main approach we decided on was to:

1. Build the model_dict mapping parameter names to real tensor values
2. Build a lazy_model_dict mapping parameter names to LazyArray values

• This will be built from the previous metadata file, which may need to be gotten through git show HEAD:path/to/checkpoint

3. Iterate through the two versions together checking lazy == real (note lazy needs to be the LHS for now)
4. Only write values when they don't evaluate to equal.

LazyArray skeleton

class LazyArray:
  def __init__(self, shape, dtype, hash):
    self._shape = shape
    self._dtype = dtype
    self._hash = hash
    
  @property
  def shape(self):
    return self._shape
    
  @property
  def dtype(self):
    return self._dtype
    
  @property
  def hash(self):
    return self._hash
    
  def __eq__(self, other):
    if self.shape != other.shape or self.dtype != other.dtype:
      return False
    if isinstance(other, LazyArray):
      return self.hash == other.hash
    else:
      return self.hash == hash_array(other)

Eventually we plan to extend the lazy array with the ability to load the value it represents.

[blester125]

Note #90 has landed. It converts the metadata file into a nested format which should make the creation of the lazy array easier

[blester125]

Note: The new multi-pointer approach #114 will change how we solve this. That change also makes this important as it allows us to skip piping parameters to git-lfs