pylibwholegraph.torch.create_embedding_from_filelist#

pylibwholegraph.torch.create_embedding_from_filelist(
comm: WholeMemoryCommunicator,
memory_type: str,
memory_location: str,
filelist: List[str] | str,
dtype: dtype,
last_dim_size: int | None = None,
*,
cache_policy: WholeMemoryCachePolicy | None = None,
embedding_entry_partition: List[int] | None = None,
gather_sms: int = -1,
round_robin_size: int = 0,
file_format: str = 'binary',
expected_entry_count: int | None = None,
expected_shape: List[int] | tuple | None = None,
fail_on_dtype_mismatch: bool = False,
)#

Create embedding from file list :param comm: WholeMemoryCommunicator :param memory_type: WholeMemory type, should be continuous, chunked or distributed :param memory_location: WholeMemory location, should be cpu or cuda :param filelist: list of binary or Parquet files :param dtype: data type :param last_dim_size: embedding width. Required for binary input and

inferred from Parquet metadata when omitted.

Parameters:
  • cache_policy – cache policy

  • embedding_entry_partition – rank partition based on entry; embedding_entry_partition[i] determines the entry count of rank i and shoud be a positive integer; the sum of embedding_entry_partition should equal to total entry count; entries will be equally partitioned if None

  • gather_sms – the number of SMs used in gather process

  • round_robin_size – continuous embedding size of a rank using round robin shard strategy

  • file_format – file format, one of binary, parquet, or auto

  • expected_entry_count – optional expected number of rows. An error is raised before allocation when the files contain a different row count.

  • expected_shape – optional expected 2-D embedding shape.

  • fail_on_dtype_mismatch – raise an error instead of warning and converting when Parquet column dtypes differ from dtype.

Returns: