Sparse Tensor: request for reviews

[eric-haibin-lin]

Following please find a summary of changes made for sparse tensor feature on CPU.

Note: this PR contains duplicate code in PR #7015

Frontend Interfaces

NDArray Class Hierarchy

-> NDArrayBase (root)
-> NDArray
-> SparseNDArray
-> RowSparseNDArray, CSRNDArray

Sparse Storage Types Overview

Two storage formats are supported: row_sparse and csr.

• row_sparse tensors, typically used for sparse gradient updates (similar to tf.IndexedSlice)

• csr tensors, with the standard csr storage format.

Functionality Overview

operators with sparse inputs

A list of operators in this PR include:

• Arithmetics: elemwise_add(row_sparse), dot(csr, row_sparse), dot(csr, dense)
• Conversion: cast_storage(csr<->dense), cast_storage(row_sparse<->dense)
• Indexing: sparse_retain(row_sparse), slice(csr)
• Creation: zeros(row_sparse), zeros(csr)

iterators and IO utils functions

• NDArrayIter with CSRNDArray
• LibsvmIter which reads files in libsvm format and produces CSRNDArray
• load/save RowSparseNDArray/CSRNDArray

kvstore and optimizer with sparse gradient updates

• sgd_update(row_sparse) and sgd_mom_update(row_sparse) will update rows of the weight where the gradient of the row contains nonzeros.
• kv.push with row_sparse gradient, which only sends over the nonzero gradients over the network
• kv.row_sparse_pull with row_ids, which only pulls back a RowSparseNDArray with specified row_id over the network

unsupported functionality

There's a list of basic functionality supported in NDArray, but is either not supported or inefficient in SparseNDArray:

• get/set an element in CSRNDArray/RowSparseNDArray by index (not supported)
• reshape in CSRNDArray/RowSparseNDArray (not supported)
• _slice, _at in CSRNDArray/RowSparseNDArray (not supported)
• inplace add/sub/mul/div/mod, broadcast_to in CSRNDArray/RowSparseNDArray (not efficient)

Please refer to the TODO list for a list of ongoing/planned operators.

Backend Interfaces

NDArray

NDArray uses storage_type to determine what storage format to use to represent a tensor.

enum NDArrayStorageType {
  kUndefinedStorage = -1,  // undefined storage
  kDefaultStorage,         // dense
  kRowSparseStorage,       // row sparse
  kCSRStorage,             // csr
};

Some of the important interfaces added to NDArray are:

NDArray {
  // constructor
  NDArray(const NDArrayStorageType stype, const TShape &shape, Context ctx,
          bool delay_alloc = true, int dtype = mshadow::default_type_flag,
          std::vector<int> aux_types = {}, std::vector<TShape> aux_shapes = {},
          TShape storage_shape = TShape(mshadow::Shape1(0)))
  // memory allocation
  void CheckAndAllocData(const TShape &storage_shape) const; 
  void CheckAndAllocAuxData(size_t i, const TShape &aux_shape) const;

  struct Chunk {
    Storage::Handle shandle;
    // storage handles for aux data (e.g. indptr, indices)
    std::vector<Storage::Handle> aux_handles;
    NDArrayStorageType storage_type;
    std::vector<int> aux_types;
    TShape storage_shape;
    std::vector<TShape> aux_shapes;
    void CheckAndAllocData(const TShape &shape, int dtype);
    void CheckAndAllocAuxData(size_t i, const TShape &shape);
  };
...

 void CopyFromTo(const NDArray &from, NDArray *to) {
    // if storage_types of from & to don't match, convert to destination storage type
    // perform copy...
 }

Executor and NNVM

operator interface

To deal with tensors of non-default storage, operators need to register with the following interface:

using FComputeEx = std::function<void (
        const nnvm::NodeAttrs& attrs,
        const OpContext& ctx,
        const std::vector<NDArray>& inputs,
        const std::vector<OpReqType>& req,
        const std::vector<NDArray>& outputs)>;

Since memory usage for sparse tensor cannot be calculated during graph binding phase, operator is expected to use CheckAndAllocAuxData and CheckAndAllocData to perform memory allocation for result tensors at runtime.

infer_storage_type pass

Besides FInferShape and FInferType, an operator needs to register FInferStorageType if it involves non-default storage_type. ctx could be used to infer if MKL experimental memory should be used. The attribute __storage_type__ of mx.sym.Variable is read for this pass.

using FInferStorageType = std::function<bool (
        const NodeAttrs& attrs,
        const Context& ctx,
        std::vector<int>* in_attrs,
        std::vector<int>* out_attrs)>;

When binding a graph, InferStorageType pass will be applied to infer graph entries with unknown storage_type. The pass is implemented in MXNet instead of NNVM. infer_type and infer_shape passes are both moved from NNVM to MXNet.

plan_memory pass

For node entries of non-default storage_type, they're annotated with kDynamicStorageID; memory are neither requested nor released. No memory sharing optimization is applied.

memory allocation and sharing

For node entries of kDefaultStorage(i.e. dense storage), memory allocation and sharing is done as usual. For node entries of non-default storage_type, memory allocation is delayed to runtime. No memory sharing is applied to these entries even when shared_exec is provided.

executor dispatch and storage fallback

For each node in graph, if it involves non-default storage in either its inputs or outputs, FComputeExExecutor will be created if FComputeEx attribute is registered.

If FComputeEx is not registered, an FComputeExcutor/FStatefulComputeExcutor will be created, inputs with non-default storage will be converted to NDArrays with default (i.e. dense) storage; the outputs will be converted back to non-default storage if necessary. The converted NDArrays are temporary and will be released when there's no reference to them.

kvstore

Sparse gradient update requires a parameter to be initialized as row_sparse format in kvstore. If an optimizer is registered with FComputeEx, then the update may only involve the rows with non-zero gradients. Gradients are sent in row_sparse format to reduce the stress on network bandwidth.

For distributed training, kRowSparsePushPull mode is introduced so that row ids of the row_sparse gradients are encoded as keys by kvstore workers. When the kvstore server receives the message, it constructs a row_sparse gradient NDArray based on the row ids decoded from it and applies the optimizer to the weight. The memory for parameters are pre-allocated on parameter servers, instead of storing weight row by row.

TODO list

Operators

• elemwise binary & unary operators (row_sparse): done, to be merged in another PR
  negative, square, abs, sqrt, round, ceil, floor, fix, elemwise_add, elemwise_sub, elemwise_div, elemwise_mul, _maximum, _minimum, _mul_scalar
• elemwise binary & unary operators (csr): ongoing
  negative, square, broadcast_mul, _mul_scalar
• inplace broadcast_add, broadcast_sub, _mul_scalar operators
• broadcast operators (csr)
• broadcast operators (row_sparse)
• elemwise operators(csr)
• _sync_copy_from with scipy.sparse.csr_matrix input
• asscipy (csr)

kvstore and optimizer

• lazy initialization - postpone parameter initialization to kvstore server
• storing row_sparse weight by rows to support weights of > 2 billion rows. This avoid pre-allocating memory for row_sparse weight and reduces memory consumption at kvstore servers.
• other optimizers for sparse gradient update (e.g. adam)

List of important files

NDArray related files:

• ndarray.h
• c_api.h
• ndarray.cc
• common/utils.h

Executor related files:

• c_api_executor.h
• attach_op_exec_pass.cc
• graph_attr_types.h
• op_attr_types.h
• executor.h
• graph_executor.h
• graph_executor.cc
• infer_graph_attr_pass.cc
• inplace_addto_detect_pass.cc

Kvstore related files:

• comm.h
• kvstore_dist.h
• kvstore_dist_server.h
• kvstore_dist_local.h
  (pending PR RowSparse pull/push #7015)

[piiswrong]

It's ok to use wild card import

[eric-haibin-lin]

will fix this in next PR

[eric-haibin-lin]

@cjolivier01 maybe you fixed it already but just in case you haven't. This bug may lead to test failure for binary/unary ops

[eric-haibin-lin]

@ZihengJiang @szha @anirudh2290 @stefanhenneking @madjam @cjolivier01
@tqchen @jermainewang @piiswrong @mli @pluskid @ykim362 @formath @reminisce

It's probably not feasible to review all the operators implemented, but I'd appreciate if important changes made to NDArray, executor and kvstore are reviewed. Any help is appreciated!

[ZihengJiang]

suggest renaming i to data_idx

[eric-haibin-lin]

@reminisce

[eric-haibin-lin]

i could either be used to indicate data or aux_data, not sure if data_idx is a good name either.

[reminisce]

Agree with @eric-haibin-lin. As long as it's documented clearly in the comment, it's fine to use i.

[ZihengJiang]

removed?

[eric-haibin-lin]

Thanks for pointing this out. will remove this comment, too

[jermainewang]

Great job! Did you consider use the storage format proposed by http://people.csail.mit.edu/fred/tensor-compiler-techreport.pdf ? It seems to me that adding more storage formats is quite tedious (e.g. CSC, BCSR, DCSR, etc.). It will be great that we can define the storage format in a hierarchical way so future extension will be much easier.

[piiswrong]

Since this cannot be rebased I'll just squash it into one commit

[fredrikbk]

Hi @eric-haibin-lin, @jermainewang, @piiswrong, all,

I'm one of the tacoers (or maybe the demonym is tacos?). I know that this pull request is closed, but I wanted to add a comment on taco. We're thinking about taco on GPUs, since it is a requirement for entry for at least ML. We think it can be done, but that it's going to be a fair amount of work to make it fast. We'll likely be working on it this spring.

There's also a preprint for our upcoming OOPSLA paper, which is much better than the tech report, at http://people.csail.mit.edu/fred/tensor-compiler-preprint.pdf.

[goswamig]

test_kvstore_gpu.test_row_sparse_pull seems to a bit flaky and has been failing randomly for egin recent following builds 475, 482. I'll create an issue for this.
Note that kv_store test runs after the "test test_forward.test_consistency" finishes.