def __init__(self,
num_numerical_features: int,
categorical_feature_sizes: Sequence[int],
bottom_mlp_sizes: Sequence[int],
top_mlp_sizes: Sequence[int],
vectors_per_gpu: Sequence[int] = None,
embedding_device_mapping: Sequence[Sequence[int]] = None,
world_num_categorical_features: int = None,
embedding_type: str = "multi_table",
embedding_dim: int = 128,
interaction_op: str = "dot",
hash_indices: bool = False,
use_cpp_mlp: bool = False,
fp16: bool = False,
bottom_features_ordered: bool = False,
device: str = "cuda"):
super().__init__()
self.distributed = dist.get_world_size() > 1
self._vectors_per_gpu = vectors_per_gpu
self._embedding_dim = embedding_dim
self._interaction_op = interaction_op
self._hash_indices = hash_indices
if self.distributed:
# TODO: take bottom_mlp GPU from device mapping, do not assume it's always first
self._device_feature_order = torch.tensor(
[-1] +
[i for bucket in embedding_device_mapping for i in bucket],
dtype=torch.long,
device=device) + 1 if bottom_features_ordered else None
self._feature_order = self._device_feature_order.argsort(
) if bottom_features_ordered else None
else:
world_num_categorical_features = len(categorical_feature_sizes)
interaction = create_interaction(interaction_op,
world_num_categorical_features,
embedding_dim)
self.bottom_model = DlrmBottom(num_numerical_features,
categorical_feature_sizes,
bottom_mlp_sizes,
embedding_type,
embedding_dim,
hash_indices=hash_indices,
use_cpp_mlp=use_cpp_mlp,
fp16=fp16,
device=device)
self.top_model = DlrmTop(top_mlp_sizes,
interaction,
use_cpp_mlp=use_cpp_mlp).to(device)
def __init__(
self,
num_numerical_features: int,
categorical_feature_sizes: Sequence[int],
bottom_mlp_sizes: Sequence[int],
top_mlp_sizes: Sequence[int],
embedding_type: str = "multi_table",
embedding_dim: int = 32,
interaction_op: str = "dot",
hash_indices: bool = False,
use_cpp_mlp: bool = False,
fp16: bool = False,
base_device: str = "cuda",
):
super().__init__()
assert embedding_dim == bottom_mlp_sizes[
-1], "The last bottom MLP layer must have same size as embedding."
interaction = create_interaction(interaction_op,
len(categorical_feature_sizes),
embedding_dim)
self._interaction_op = interaction_op
self._hash_indices = hash_indices
self.bottom_model = DlrmBottom(
num_numerical_features=num_numerical_features,
categorical_feature_sizes=categorical_feature_sizes,
bottom_mlp_sizes=bottom_mlp_sizes,
embedding_type=embedding_type,
embedding_dim=embedding_dim,
hash_indices=hash_indices,
use_cpp_mlp=use_cpp_mlp,
fp16=fp16,
device=base_device)
self.top_model = DlrmTop(top_mlp_sizes,
interaction,
use_cpp_mlp=use_cpp_mlp).to(base_device)