mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
16549ed92b
Summary: Today, the PS's weirdly store the entire embedding and not just their subsection of it. This was simply an oversight on the part of the original author and this diff fixes that. The sparse params are sharded to the PS's and the PS's just store their section of the embedding. The trainer requests the id's as is from the PS. But the PS divides the id by the num_of_shards before looking it up in the emdedding table blob. This happens on the backward and the forward pass. However, during the model download part, the PS multiples the embeddings with the num_of_shards before returning them to the trainer. The upshot is that the trainer does not know anything about how the embeddings are scaled on the PS. The PS adds extra divide and multiply steps to achieve that. 2. During estimation time, we allocate just one PS for estimation. So in order to make all of the embeddings fit on the single PS: We simply additionally scale the hash table sizes (proportionally and equally for all the sparse params) such that it fits. This scaling is handled analogously to (1). Reviewed By: boryiingsu Differential Revision: D5664093 fbshipit-source-id: 92f501f61566f939c41ce0b614a1b499669f978a
68 lines
1.9 KiB
Python
68 lines
1.9 KiB
Python
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
from __future__ import unicode_literals
|
|
|
|
from caffe2.python import core
|
|
|
|
import numpy as np
|
|
|
|
|
|
class ParameterTags(object):
|
|
BIAS = 'BIAS'
|
|
WEIGHT = 'WEIGHT'
|
|
COMPUTED_PARAM = 'COMPUTED_PARAM'
|
|
|
|
|
|
class ParameterType(object):
|
|
DENSE = 'dense'
|
|
SPARSE = 'sparse'
|
|
|
|
|
|
class ParameterInfo(object):
|
|
|
|
def __init__(
|
|
self, param_id, param, key=None, shape=None, length=None,
|
|
grad=None, blob_copy=None):
|
|
assert isinstance(param, core.BlobReference)
|
|
self.param_id = param_id
|
|
self.name = str(param)
|
|
self.blob = param
|
|
self.key = key
|
|
self.shape = shape
|
|
self.size = None if shape is None else np.prod(shape)
|
|
self.length = max(1, length if length is not None else 1)
|
|
self.grad = grad
|
|
self._cloned_init_net = None
|
|
# Optionally store equivalent copies of the blob
|
|
# in different precisions (i.e. half and float copies)
|
|
# stored as a dict of TensorProto.DataType -> BlobReference
|
|
self.blob_copy = blob_copy
|
|
# each param_info can have its own optimizer. It can be set within
|
|
# OptimizerContext (caffe2/python/optimizer.py)
|
|
self._optimizer = None
|
|
|
|
def grad_type(self):
|
|
# self.grad could be None for model parallelism with parameter server
|
|
if self.grad is None:
|
|
return
|
|
return (
|
|
ParameterType.SPARSE if isinstance(self.grad, core.GradientSlice)
|
|
else ParameterType.DENSE)
|
|
|
|
@property
|
|
def parameter(self):
|
|
return self.blob
|
|
|
|
@property
|
|
def optimizer(self):
|
|
return self._optimizer
|
|
|
|
@optimizer.setter
|
|
def optimizer(self, value):
|
|
assert self._optimizer is None, "optimizer has already been set"
|
|
self._optimizer = value
|
|
|
|
def __str__(self):
|
|
return self.name
|