JasonMoho
commented
1 year ago It definitely looks like you should have enough memory to preprocess this dataset. Let's get to the bottom of this.
The OOM error is on an array with 167 million elements and takes up 612GiB. So this means that each element (node_id) takes up about ~4KB on average. That seems far too high for a simple node_id, especially since your triple file only has 46GB (which implies only 46 bytes per triple with 1B triples).
What does the tsv look like? Worth checking if the delimiter is actually a tab. Does it look something like:
src_label rel_label dst_label
"node_src_0" "rel_0" "node_dst_0"
"node_src_1" "rel_1" "node_dst_1"
"node_src_2" "rel_2" "node_dst_2"
"node_src_3" "rel_3" "node_dst_3"I'm wondering if numpy is doing something strange in the backend when handling strings, like allocating too much memory per element in this operation: np.concatenate([unique_src.astype(str), unique_dst.astype(str)])
We have two options for a workaround here:
- Write a python script to call the
TorchEdgeListConvertersimilar to ogbl_citation2. Where a 2D numpy array for your edge list is passed into theTorchEdgeListConverterclass, and the raw string ids are converted into integer ids manually before input. Usingnp.loadtxt("/mount/data/works.tsv")with the proper arguments should work to read in your edge list. This will go down a different code path than the line that throws the error and should bypass the problem.
The usage of the TorchEdgeListConverter will look like:
import numpy as np
from marius.tools.preprocess.converters.torch_converter import TorchEdgeListConverter
train_list = np.loadtxt("/mount/data/works.tsv", delimiter="/t", skiprows=header_length)
# convert string ids to integer ids using custom code written by you
train_list = convert_to_integers(train_list)
converter = TorchEdgeListConverter(
output_dir="/mount/marius_preprocessed/works/",
train_edges=train_list,
num_partitions=4,
columns=[0, 1, 2],
format="numpy",
splits=[.8, .1, .1])
converter.convert()Full list of arguments for the TorchEdgeListConverter located here
- Try using spark preprocessing. This requires pyspark (and the jvm) installed to run. You can use this version of the preprocessing by passing
--sparkto the preprocessor. This will be a bit slower than the in-memory preprocessing but should be able to handle strings better.
johankit
Problem description
When preprocessing my data set in the .tsv format to prepare it for training and splitting into four partitions, I receive an out-of-memory error.
I used the command:
marius_preprocess --output_dir /mount/marius_preprocessed/works/ --edges /mount/data/works.tsv --dataset_split 0.8 0.1 0.1 --columns 0 1 2 --num_partitions 4However, during preprocessing I encounter the error:
unique_nodes = np.unique(np.concatenate([unique_src.astype(str), unique_dst.astype(str)])) numpy.core._exceptions.MemoryError: Unable to allocate 612GiB for an array with shape (167416627, ).The input file is 46GB in size and contains about 1 billion lines (triples). And the instance I'm using has 500GB of memory. It seems the array has the 'length' equal to the number of unique entities in the input file.
The error occurs after the remapping of edges step has started. Changing the number of partitions did not help. I am running the tool in a Docker container, the behavior without container was similar though.
I understand Marius was built for efficient embeddings generation on lower-capacity machines when training. Is there any way to reduce the resource needs during preproccesing as well? Perhaps any modifications in the .tsv file from my side that could support the preprocessing?
Expected behavior
Preprocessing of raw input (.nt or .tsv files) into ready-to-train partitions, with comparable resource requirements as during the embeddings training.
Environment
Thank you!