[GitHub] eric-haibin-lin opened a new issue #9921: [DISCUSSION] module.contrib.SparseModule API

[GitBox <gi...@apache.org>]

eric-haibin-lin opened a new issue #9921: [DISCUSSION] module.contrib.SparseModule API
URL: https://github.com/apache/incubator-mxnet/issues/9921
 
 
   As of now, the only way to train a model with sparse operators is using the Module API (sparse support for Gluon API not ready yet). However, the existing Module API is not a good fit for training sparse models with KVStores present. This thread discusses how to provide a user-friendly module API for sparse symbols, which serves as a guideline before the feature is added into Gluon.
   
   ## Why the current Module API is not a good fit
   Currently, if a module has KVStore present, it performs the following steps for each batch:
   ```
   for batch in batches:
       mod.forward(batch)
       mod.backward(batch)
       mod.update() which calls the following kvstore APIs
           # kvstore.push(grad) which aggregates & applies the gradients from the current batch
           # kvstore.pull(weight) which broadcast updated weights to all devices
   ```
   This is not ideal for models w/ sparse operators because it is not necessary to broadcast full weights. Instead, only a subset of the weights are useful for the computation when data is sparse. 
   
   Here is an example:
   Let's say input `X` is has 100,000 features, where each sample only has around 10 non-zeros features. If the model computes `dot(X, W)`, not all rows in `W` are used for the computation. In fact, if `W` only contains the rows corresponding to the non-zero features in `X`, the result of `dot` is exactly the same compared to a `dot(X, W)` w/ full(dense) `W`. 
   
   Avoiding the communication of these full(dense) weight/gradient is essential when one tries to scale the training over multiple machines and multiple devices. 
   
   Therefore, we want to re-order the loop and do something like the following:
   ```
   for batch in batches:
       kvstore.row_sparse_pull(sparse_weight, useful_rows_for_this_batch) 
       kvstore.pull(dense_weights)
       mod.forward(batch)
       mod.backward(batch) 
       mod.update()
           # kvstore.push(grad) 
   ```
   With the module API, one must access private members of module such as `_exec_group` and `_params_dirty` and interact with kvstore directly to achieve the above like this [example](https://github.com/eric-haibin-lin/mxnet/blob/lstmp/example/sparse/linear_classification/train.py#L97-L98). We need a module API that abstracts these away and make it more user-friendly. 
   
   ## contrib.SparseModule API
   Modifying the existing module API to support this may over complicate the code and make it hard to maintain. On the other hand, a new SparseModule API is a chance to educate MXNet users who already have some experience with the Module API. I propose the following SparseModule API: 
   
   ### class SparseModule(Module)
   
       A Module that wraps a network with sparse symbols. 
       The module provides abstraction for data parallel distributed training with "row sparse" weights,
       supporting two types of communications: update_on_kvstore=True and update_on_kvstore=False.
       Distributed training includes both multi-device and multi-machine training.
   
        - update_on_kvstore=True
        Full parameters are stored in KVStore.
        For each data batch, only broadcast the subset of "row sparse" weights that are required for the computation,
        according to the row ids provided. 
        Usually the row ids correspond to features that appear in the batch.
   
        Specifically, for each batch of data, 
        sparse weights are broadcast to all devices/machines based on row ids. 
        Each device/machine then computes forward/backward with sparse weights.
        Finally, gradients are reduced from all devices/machines and applied on the parameters stored in KVStore.
   
        - update_on_kvstore=False
        Full parameters are stored and duplicated on each machine/device. 
        For each data batch, compute forward/backward with full weights.
        Then gradients are then broadcast to all other devices/machines. 
        Each device/machine applies the gradient to the copy of the parameters it holds. 
   
       Take the following linear network as an example:
       ```python
        data = mx.sym.var('data', stype='csr')
        label = mx.sym.var('label')
        weight = mx.sym.var('w0', shape=(1000,1), stype='row_sparse')
        bias = mx.sym.var('b0', shape=(1,))
        pred = mx.sym.broadcast_add(mx.sym.sparse.dot(data, weight), bias)
        output = mx.sym.LogisticRegressionOutput(pred, label)
        ```
   
        - update_on_kvstore=True
        ```python
        mod = mx.module.contrib.SparseModule(..., update_on_kvstore=True)
        for batch in batches:
            row_id_generator = lambda x: {'w0': x.data[0].indices}
            **mod.prepare(batch, row_id_generator=row_id_generator)**
            mod.forward(batch)
            mod.backward()
            mod.update()
        ```
   
        - update_on_kvstore=False
        ```python
        mod = mx.module.contrib.SparseModule(..., update_on_kvstore=False)
        for batch in batches:
            **mod.prepare(batch)**
            mod.forward(batch)
            mod.backward()
            mod.update()
        ```
   
   #### __init__(..., **update_on_kvstore=None**)
       update_on_kvstore : bool 
           Specifies whether to store a copy of the parameter and perform updates on the KVStore. 
           The value type could be extended to "bool or dict" in the future,
           if we want to support different behavior per parameter, 
           where the keys of the dictionary are the names of parameters, the values are boolean.
           e.g. {'weight_0':False, 'weight_1':True}.
   
   #### prepare(batch, row_id_generator=None)
       Prepare the module for the current data batch. 
       update_on_kvstore=True, their updated values are pulled from the KVStore,
       based on the row ids generated by row_id_generator.
   
       row_id_generator: a lambda function which takes data_batch as the input, 
             and outputs a dictionary of str->NDArray. The keys in the dict are sparse param names,
             the values are the row_ids to pull for each param.
   
   #### update
       Update parameters based on the gradients computed by the previous forward/backward.
       If update_on_kvstore=True, it updates the parameters stored in KVStore.
       Otherwise it updates the params stored locally on each device. 
   
   #### init_optimizer
       initialize KVStore and optimizer. 
       There is no guarantee that the parameters on devices are in sync with the ones in KVStore. 
       The sync happens in `prepare()` instead. 
   
   #### save_checkpoint
       if update_on_kvstore=True, need to pull sparse weights before saving the checkpoint.
   
   #### score
       if update_on_kvstore=True, need to pull sparse weights before calling forward. 
   
   #### iter_predict
       if update_on_kvstore=True, need to pull sparse weights before calling forward. 
   
   #### predict
       if update_on_kvstore=True, need to pull sparse weights before calling forward. 
   
   #### fit
       if update_on_kvstore=True, need to pull sparse weights before calling forward. 
   
   #### get_params
       Gets a copy of current parameters.
       If update_on_kvstore=True, need to pull all sparse weights. 
   
   #### save_params
       if update_on_kvstore=True, need to pull all sparse weights. 
   
   #### forward_backward
   
   #### @property data_names
   #### @property output_names
   #### @property data_shapes
   #### @property label_shapes
   #### output_shapes
   #### init_params
   #### set_params
   
   #### load_params
   #### get_states
   #### set_states
   #### install_monitor
   
   #### forward
   #### backward
   #### get_outputs
   #### get_input_grads
   #### update_metric
   #### bind
   
   #### @property symbol
   #### @staticmethod load
   
   #### @property label_names
   #### reshape
   #### borrow_optimizer
   #### save_optimizer_states
   #### load_optimizer_states
   

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services