{"action":"created","issue":{"url":"https://api.github.com/repos/apache/incubator-mxnet/issues/9921","repository_url":"https://api.github.com/repos/apache/incubator-mxnet","labels_url":"https://api.github.com/repos/apache/incubator-mxnet/issues/9921/labels{/name}","comments_url":"https://api.github.com/repos/apache/incubator-mxnet/issues/9921/comments","events_url":"https://api.github.com/repos/apache/incubator-mxnet/issues/9921/events","html_url":"https://github.com/apache/incubator-mxnet/issues/9921","id":301014622,"number":9921,"title":"[DISCUSSION] module.contrib.SparseModule API","user":{"login":"eric-haibin-lin","id":5545640,"avatar_url":"https://avatars1.githubusercontent.com/u/5545640?v=4","gravatar_id":"","url":"https://api.github.com/users/eric-haibin-lin","html_url":"https://github.com/eric-haibin-lin","followers_url":"https://api.github.com/users/eric-haibin-lin/followers","following_url":"https://api.github.com/users/eric-haibin-lin/following{/other_user}","gists_url":"https://api.github.com/users/eric-haibin-lin/gists{/gist_id}","starred_url":"https://api.github.com/users/eric-haibin-lin/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/eric-haibin-lin/subscriptions","organizations_url":"https://api.github.com/users/eric-haibin-lin/orgs","repos_url":"https://api.github.com/users/eric-haibin-lin/repos","events_url":"https://api.github.com/users/eric-haibin-lin/events{/privacy}","received_events_url":"https://api.github.com/users/eric-haibin-lin/received_events","type":"User","site_admin":false},"labels":[{"id":272547979,"url":"https://api.github.com/repos/apache/incubator-mxnet/labels/Discussion","name":"Discussion","color":"0e8a16","default":false},{"id":770413511,"url":"https://api.github.com/repos/apache/incubator-mxnet/labels/Module","name":"Module","color":"fbca04","default":false},{"id":720179133,"url":"https://api.github.com/repos/apache/incubator-mxnet/labels/Sparse","name":"Sparse","color":"fbca04","default":false}],"state":"open","locked":false,"assignee":null,"assignees":[],"milestone":null,"comments":8,"created_at":"2018-02-28T12:59:18Z","updated_at":"2018-03-09T16:00:32Z","closed_at":null,"author_association":"CONTRIBUTOR","body":"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.\r\n\r\n## Why the current Module API is not a good fit\r\nCurrently, if a module has KVStore present, it performs the following steps for each batch:\r\n```python\r\nfor batch in batches:\r\n    mod.forward(batch)\r\n    mod.backward(batch)\r\n    mod.update() which calls the following kvstore APIs\r\n        # kvstore.push(grad) which aggregates & applies the gradients from the current batch\r\n        # kvstore.pull(weight) which broadcast updated weights to all devices\r\n```\r\nThis 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. \r\n\r\nHere is an example:\r\nLet'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`. \r\n\r\nAvoiding the communication of these full(dense) weight/gradient is essential when one tries to scale the training over multiple machines and multiple devices. \r\n\r\nTherefore, we want to re-order the loop and do something like the following:\r\n```python\r\nfor batch in batches:\r\n    kvstore.row_sparse_pull(sparse_weight, useful_rows_for_this_batch) \r\n    kvstore.pull(dense_weights)\r\n    mod.forward(batch)\r\n    mod.backward(batch) \r\n    mod.update()\r\n        # kvstore.push(grad) \r\n```\r\nWith 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. \r\n\r\n## contrib.SparseModule API\r\nModifying 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: \r\n\r\n### class SparseModule(Module)\r\n\r\nA Module that wraps a network with sparse symbols.\r\nThe module provides abstraction for data parallel distributed training with \"row sparse\" weights,\r\nsupporting two types of communications: update_on_kvstore=True and update_on_kvstore=False.\r\nDistributed training includes both multi-device and multi-machine training.\r\n\r\n- update_on_kvstore=True\r\n    Full parameters are stored in KVStore.\r\n    For each data batch, only broadcast the subset of \"row sparse\" weights that are required for the computation,\r\n    according to the row ids provided.\r\n    Usually the row ids correspond to features that appear in the batch.\r\n    \r\n    Specifically, for each batch of data,\r\n    sparse weights are broadcast to all devices/machines based on row ids.\r\n    Each device/machine then computes forward/backward with sparse weights.\r\n    Finally, gradients are reduced from all devices/machines and applied on the parameters stored in KVStore.\r\n\r\n\r\n- update_on_kvstore=False\r\n    Full parameters are stored and duplicated on each machine/device.\r\n    For each data batch, compute forward/backward with full weights.\r\n    Then gradients are then broadcast to all other devices/machines.\r\n    Each device/machine applies the gradient to the copy of the parameters it holds.\r\n\r\nTake the following linear network as an example:\r\n```python\r\n data = mx.sym.var('data', stype='csr')\r\n label = mx.sym.var('label')\r\n weight = mx.sym.var('w0', shape=(1000,1), stype='row_sparse')\r\n bias = mx.sym.var('b0', shape=(1,))\r\n pred = mx.sym.broadcast_add(mx.sym.sparse.dot(data, weight), bias)\r\n output = mx.sym.LogisticRegressionOutput(pred, label)\r\n ```\r\n\r\n - update_on_kvstore=True\r\n ```python\r\n mod = mx.module.contrib.SparseModule(..., update_on_kvstore=True)\r\n for batch in batches:\r\n     row_id_generator = lambda x: {'w0': x.data[0].indices}\r\n     **mod.prepare(batch, row_id_generator=row_id_generator)**\r\n     mod.forward(batch)\r\n     mod.backward()\r\n     mod.update()\r\n ```\r\n\r\n - update_on_kvstore=False\r\n ```python\r\n mod = mx.module.contrib.SparseModule(..., update_on_kvstore=False)\r\n for batch in batches:\r\n     **mod.prepare(batch)**\r\n     mod.forward(batch)\r\n     mod.backward()\r\n     mod.update()\r\n ```\r\n\r\nWhich type is recommended? So far only `update_on_kvstore=True` is implemented.\r\n- Memory: `update_on_kvstore=False` in theory leads to higher memory consumption,\r\n   because each device keeps a complete copy of the model.\r\n   For extremely large models that don't fit on a single machine/device,\r\n   keeping a complete copy of the model is not feasible.\r\n- Computation: depends on the operator. For some operators,\r\n   they benefit from the fact that weights are sparse (such as square_sum in factorization machine)\r\n   and do less computation during forward.\r\n   For other operators it might not be the case\r\n   because sparse weight requires extra efforts for lookup/search(such as SparseEmbedding).\r\n\r\n#### __init__(..., **update_on_kvstore=None**)\r\n    update_on_kvstore : bool \r\n        Specifies whether to store a copy of the parameter and perform updates on the KVStore. \r\n        The value type could be extended to \"bool or dict\" in the future,\r\n        if we want to support different behavior per parameter, \r\n        where the keys of the dictionary are the names of parameters, the values are boolean.\r\n        e.g. {'weight_0':False, 'weight_1':True}.\r\n\r\n#### prepare(batch, row_id_generator=None)\r\n    Prepare the module for the current data batch. \r\n    update_on_kvstore=True, their updated values are pulled from the KVStore,\r\n    based on the row ids generated by row_id_generator.\r\n\r\n    row_id_generator: a lambda function which takes data_batch as the input, \r\n          and outputs a dictionary of str->NDArray. The keys in the dict are sparse param names,\r\n          the values are the row_ids to pull for each param.\r\n\r\n    Note that the reason why row_id_generator is a function is that users can customize what row ids to provide. It could be the `indices` field from a CSRNDArray, or the input data itself (such as user_id, movie_id). \r\n\r\n#### update\r\n    Update parameters based on the gradients computed by the previous forward/backward.\r\n    If update_on_kvstore=True, it updates the parameters stored in KVStore.\r\n    Otherwise it updates the params stored locally on each device. \r\n\r\n#### init_optimizer\r\n    initialize KVStore and optimizer. \r\n    There is no guarantee that the parameters on devices are in sync with the ones in KVStore. \r\n    The sync happens in `prepare()` instead. \r\n\r\n#### save_checkpoint\r\n    if update_on_kvstore=True, need to pull sparse weights before saving the checkpoint.\r\n\r\n#### score\r\n    if update_on_kvstore=True, need to pull sparse weights before calling forward. \r\n\r\n#### iter_predict\r\n    if update_on_kvstore=True, need to pull sparse weights before calling forward. \r\n\r\n#### predict\r\n    if update_on_kvstore=True, need to pull sparse weights before calling forward. \r\n\r\n#### fit\r\n    if update_on_kvstore=True, need to pull sparse weights before calling forward. \r\n\r\n#### get_params\r\n    Gets a copy of current parameters.\r\n    If update_on_kvstore=True, need to pull all sparse weights. \r\n\r\n#### save_params\r\n    if update_on_kvstore=True, need to pull all sparse weights. \r\n\r\n#### forward_backward\r\n\r\n#### @property data_names\r\n#### @property output_names\r\n#### @property data_shapes\r\n#### @property label_shapes\r\n#### output_shapes\r\n#### init_params\r\n#### set_params\r\n\r\n#### load_params\r\n#### get_states\r\n#### set_states\r\n#### install_monitor\r\n\r\n#### forward\r\n#### backward\r\n#### get_outputs\r\n#### get_input_grads\r\n#### update_metric\r\n#### bind\r\n\r\n#### @property symbol\r\n#### @staticmethod load\r\n\r\n#### @property label_names\r\n#### reshape\r\n#### borrow_optimizer\r\n#### save_optimizer_states\r\n#### load_optimizer_states\r\n"},"comment":{"url":"https://api.github.com/repos/apache/incubator-mxnet/issues/comments/371854781","html_url":"https://github.com/apache/incubator-mxnet/issues/9921#issuecomment-371854781","issue_url":"https://api.github.com/repos/apache/incubator-mxnet/issues/9921","id":371854781,"user":{"login":"zheng-da","id":70481,"avatar_url":"https://avatars1.githubusercontent.com/u/70481?v=4","gravatar_id":"","url":"https://api.github.com/users/zheng-da","html_url":"https://github.com/zheng-da","followers_url":"https://api.github.com/users/zheng-da/followers","following_url":"https://api.github.com/users/zheng-da/following{/other_user}","gists_url":"https://api.github.com/users/zheng-da/gists{/gist_id}","starred_url":"https://api.github.com/users/zheng-da/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/zheng-da/subscriptions","organizations_url":"https://api.github.com/users/zheng-da/orgs","repos_url":"https://api.github.com/users/zheng-da/repos","events_url":"https://api.github.com/users/zheng-da/events{/privacy}","received_events_url":"https://api.github.com/users/zheng-da/received_events","type":"User","site_admin":false},"created_at":"2018-03-09T16:00:32Z","updated_at":"2018-03-09T16:00:32Z","author_association":"CONTRIBUTOR","body":"@eric-haibin-lin True. We don't need to have graph executor or operators aware of kvstore. What we need to do is to register a function to NNVM for operators. What the function does is to return row ids for given input. The process is kind of similar to shape inference but is called when the inputs are ready. The graph executor calls these functions (if they are available) and accumulates all row ids required for the computation. Some component that is responsible for communicating with kvstore should send these row Ids to kvstore and fetch weights.\r\n\r\nWhat I mean is to use an array to map row ids to the physical location of a row in the row sparse matrix, instead of using a hashtable. My point is that if you want to use `update_on_kvstore=False`, which costs memory of O(n * p), why not use an array to map row ids, which costs memory of O(n)?  It isn't expensive to maintain the map either, in terms of computation. The cost is O(nzr) if you do it correctly. I think maintaining this array map is strictly better than maintaining the entire weight matrix in every worker node, from the perspective of computation and memory overhead."}}
{"actor":{"display_login":"zheng-da"}}