rewrite the code of update and broadcast workers

source/train/XLeader.cpp

@@ -267,6 +267,11 @@ run the model (for one time). Basically this is a map-reduce process.
 bool XLeader::Run(XConfig * config, DataDistributeBase * dataDistributor,
                   XModel * model, XOptimizer * optimizer)
 {
+    CheckNTErrors(jworkers.count > 0, "No jworkers!");
+    CheckNTErrors(cworkers.count > 0, "No cworkers!");
+    CheckNTErrors(uworkers.count > 0, "No uworkers!");
+    CheckNTErrors(bworkers.count > 0, "No bworkers!");
+
     bool isDataOK = true;
     int activeJobCount = 0;
     int* active = new int[jworkers.count];
@@ -306,7 +311,12 @@ bool XLeader::Run(XConfig * config, DataDistributeBase * dataDistributor,
         }
     }
 
-    if (activeJobCount >= 0) {
+    if (activeJobCount > 0) {
+        /* workers */
+        XWorkerCollect * collecter = (XWorkerCollect*)cworkers.GetItem(0);
+        XWorkerUpdate * updater = (XWorkerUpdate*)uworkers.GetItem(0);
+        XWorkerBroadcast * broadcaster = (XWorkerBroadcast*)bworkers.GetItem(0);
+
         /* member models that are active in this run */
         XList members(jworkers.count);
 
@@ -325,36 +335,22 @@ bool XLeader::Run(XConfig * config, DataDistributeBase * dataDistributor,
             }
         }
 
+        collecter->AddJobUpdateAll(&members, &membersAll, &serverModel, 
+                                   optimizer, updater, broadcaster);
+        collecter->AddJobCollectOther(&memberRecords, &serverRecord);
+
         /* jobs in queue 2: collect the (gradient) data and other stuff. This
            is a reduce process. */
-        if (cworkers.count > 0) {
-            XWorkerCollect* collecter = (XWorkerCollect*)cworkers.GetItem(0);
-            collecter->AddJobCollect(&members, &serverModel);
-            collecter->AddJobCollectOther(&memberRecords, &serverRecord);
-        }
-        else {
-            ShowNTErrors("No data-collecting workers!");
-        }
+        //collecter->AddJobCollect(&members, &serverModel);
+        //collecter->AddJobCollectOther(&memberRecords, &serverRecord);
 
         /* job in queue 3: update the model */
-        if (uworkers.count > 0) {
-            XWorkerUpdate* updater = (XWorkerUpdate*)uworkers.GetItem(0);
-            updater->AddJobUpdate(&serverModel, optimizer);
-        }
-        else {
-            ShowNTErrors("No model-update workers!");
-        }
+        //updater->AddJobUpdate(&serverModel, optimizer);
 
         /* job in queue 4: broadcast the lastest parameters to workers. NOTE that
            we would update a worker to the laster model parameters, even if it is
            not involved in this run. */
-        if (bworkers.count > 0) {
-            XWorkerBroadcast* broadcaster = (XWorkerBroadcast*)bworkers.GetItem(0);
-            broadcaster->AddJobBroadcast(&serverModel, &membersAll);
-        }
-        else {
-            ShowNTErrors("No data-broadcasting workers!");
-        }
+        //broadcaster->AddJobBroadcast(&serverModel, &membersAll);
 
         WaitForFinishing();
     }

source/train/XWorkerBroadcast.cpp

@@ -50,7 +50,32 @@ void XWorkerBroadcast::SetBroadcastMode(DATA_BROADCAST_TYPE myMode)
 }
 
 /* 
-broadcast data 
+broadcast data for a parameter 
+>> source - the data (as a model) that we want to broadcast
+>> targetList - the target places that we recieve the data
+>> pid - the parameter index
+*/
+void XWorkerBroadcast::BroadcastDataSingle(XModel * source, XList * targetList, int pid)
+{
+    CheckNTErrors(source->flags[pid] == PARAM_STATE_UPDATED,
+                  "The parameter is not ready for broadcasting");
+
+    TensorList & sp = source->params;
+
+    for (int i = 0; i < targetList->count; i++) {
+        XModel * target = (XModel*)targetList->GetItem(i);
+        TensorList & tp = target->params;
+
+        /* data transmit */
+        BroadcastP2P(sp.GetItem(pid), tp.GetItem(pid));
+
+        /* update the flag */
+        target->flags[pid] = PARAM_STATE_UPDATED;
+    }
+}
+
+/* 
+broadcast data for a model
 >> source - the data that we want to broadcast
 >> targetList - the target places that we recieve the data
 >> sleepTime - the waiting time in broadcasting
@@ -72,17 +97,12 @@ void XWorkerBroadcast::BroadcastData(XModel * source, XList * targetList, long s
     while (1) {
         for (int i = 0; i < sp.count; i++) {
             if (source->flags[i] == PARAM_STATE_UPDATED && finishedFlag[i] == 0) {
-                for (int j = 0; j < targetList->count; j++) {
-                    XModel * target = (XModel*)targetList->GetItem(j);
-                    TensorList & tp = target->params;
 
-                    /* data transmit */
-                    BroadcastP2P(sp.GetItem(i), tp.GetItem(i));
+                /* broadcasting */
+                BroadcastDataSingle(source, targetList, i);
 
-                    /* update the flag */
-                    target->flags[i] = PARAM_STATE_UPDATED;
-                    finished++;
-                }
+                /* counting */
+                finished += targetList->count;
                 finishedFlag[i] = 1;
             }
         }
@@ -97,6 +117,29 @@ void XWorkerBroadcast::BroadcastData(XModel * source, XList * targetList, long s
 }
 
 /* 
+wrapper of BroadcastDataSingle 
+>> args - the list of arguments
+*/
+void XWorkerBroadcast::BroadcastSingle(XList * args)
+{
+    XWorkerBroadcast * broadcaster = (XWorkerBroadcast*)args->GetItem(0);
+    XModel * source = (XModel*)args->GetItem(1);
+
+    /* target models */
+    int targetNum = args->GetItemInt(2);
+    XList target;
+    for (int i = 0; i < targetNum; i++) {
+        XModel * model = (XModel*)args->GetItem(3 + i);
+        target.Add(model);
+    }
+
+    /* parameter index */
+    int p = args->GetInt(3 + targetNum);
+
+    broadcaster->BroadcastDataSingle(source, &target, p);
+}
+
+/* 
 wrapper of BroadcastData 
 >> args - the list of arguments
 */
@@ -129,11 +172,39 @@ void XWorkerBroadcast::BroadcastP2P(XTensor * source, XTensor * target)
     CheckNTErrors(target != NULL, "The target tensor should not be NULL!");
     CheckNTErrors(IsSameShaped(*source, *target), "The two tensors should be of the same shape!");
 
+    if(source != target)
         CopyValues(*source, *target);
 }
 
 /* 
-add a new job of broadcasting data 
+add a new job of broadcasting data (for a parameter)
+>> source - the data that we want to broadcast
+>> targetList - the target places that we recieve the data
+>> pid - the parameter index
+*/
+bool XWorkerBroadcast::AddJobBroadcastSingle(XModel * source, XList * targetList, int pid)
+{
+    CheckNTErrors(source != NULL, "no input source tensor!");
+    CheckNTErrors(targetList != NULL, "no input target tensor list!");
+    CheckNTErrors(pid >= 0 && pid < source->params.count, "illegal parameter index!");
+
+    XList args;
+    args.Add(this);
+    args.Add(source);
+    args.AddInt(targetList->count);
+    args.AddList(targetList);
+    args.AddInt(pid);
+
+    if (isInstantRun)
+        XWorkerBroadcast::BroadcastSingle(&args);
+    else
+        queue.EnqueueJob((void*)(char*)XWorkerBroadcast::BroadcastSingle, &args);
+
+    return true;
+}
+
+/* 
+add a new job of broadcasting data (for a model)
 >> source - the data that we want to broadcast
 >> targetList - the target places that we recieve the data
 */

source/train/XWorkerBroadcast.h

@@ -60,9 +60,16 @@ public:
     /* set the broadcasting type */
     void SetBroadcastMode(DATA_BROADCAST_TYPE myMode);
 
-    /* broadcast data */
+    /* broadcast data for a parameter */
+    void BroadcastDataSingle(XModel * source, XList * targetList, int pid);
+
+    /* broadcast data for a model */
     void BroadcastData(XModel * source, XList * targetList, long sleepTime);
 
+    /* wrapper of BroadcastDataSingle */
+    static
+    void BroadcastSingle(XList * args);
+
     /* wrapper of BroadcastData */
     static
     void Broadcast(XList * args);
@@ -70,7 +77,10 @@ public:
     /* P2P data broadcasting */
     void BroadcastP2P(XTensor * source, XTensor * target);
 
-    /* add a new job of broadcasting data */
+    /* add a new job of broadcasting data (for a parameter) */
+    bool AddJobBroadcastSingle(XModel * source, XList * targetList, int pid);
+
+    /* add a new job of broadcasting data (for a model) */
     bool AddJobBroadcast(XModel * source, XList * targetList);
 };
 

source/train/XWorkerCollect.cpp

@@ -49,6 +49,165 @@ void XWorkerCollect::SetCollectMode(DATA_COLLECT_TYPE myMode)
 }
 
 /* 
+collect the gradient data, update the parameters, and broadcast the
+new parameters to all models. NOTE that this method just collect graident
+from member models. Then it calls an XWorkerUpdate to update the parameters.
+The XWorkerUpdate also calls an XWorkerBroadcast to broadcast the new parameter
+to member models back. 
+>> memberActive - member models that are active, i.e., have generated gradients
+>> memberAll -  all member models
+>> server - the server model
+>> optimizer - the optimizer
+>> updater - the worker that updates the parameters
+>> broadcaster - the worker that broadcasts the new parameters to all member
+                 models
+>> sleepTime - waiting time in collecting
+*/
+void XWorkerCollect::UpdateDataAll(XList * memberActive, XList * memberAll, XModel * server,
+                                   XOptimizer * optimizer, XWorkerUpdate * updater, 
+                                   XWorkerBroadcast * broadcaster, long sleepTime)
+{
+    TensorList & tp = server->params;
+    int finished = 0;
+
+    for (int j = 0; j < tp.count; j++)
+        server->flags[j] = PARAM_STATE_NOT_READY;
+
+    /* check */
+    for (int i = 0; i < memberAll->count; i++) {
+        TensorList & sp = ((XModel*)memberAll->GetItem(i))->params;
+        CheckNTErrors(sp.count == tp.count, "Incompatiable models!");
+    }
+
+    for (int i = 0; i < memberActive->count; i++) {
+        TensorList & sp = ((XModel*)memberActive->GetItem(i))->params;
+        CheckNTErrors(sp.count == tp.count, "Incompatiable models!");
+    }
+
+    /* counts how many member models are collect for each parameters */
+    int * finishedCount = new int[tp.count];
+    memset(finishedCount, 0, sizeof(int) * tp.count);
+
+    /* This is a simple implementation of the wait-and-collect process. But
+       there is a risk that some models are not available, that is, the
+       loop would never stop. A solution might be that we force the loop
+       to break after waiting for a short time. */
+    while (1) {
+        if (collectMode == DATA_COLLECT_P2P) {
+            for (int j = 0; j < tp.count; j++) {
+                /* tp[j]->isGradFinished is true only if the model finishes the computation
+                (in another process) */
+                if (server->flags[j] != PARAM_STATE_NOT_READY || !tp[j]->isGradFinished)
+                    continue;
+
+                /* check if all the models (or part of them) are ready */
+                for (int i = 0; i < memberActive->count; i++) {
+                    XModel * source = (XModel*)memberActive->GetItem(i);
+                    TensorList & sp = source->params;
+
+                    /* sp[j]->isGradFinished is true only if the model finishes the computation
+                    (in another process) */
+                    if (source->flags[j] == PARAM_STATE_NOT_READY && sp[j]->isGradFinished) {
+
+                        /* data transmit */
+                        CollectP2P(sp[j]->grad, tp[j]->grad);
+
+                        /* reset the flag */
+                        source->flags[j] = PARAM_STATE_COLLECTED;
+                        finished++;
+                        finishedCount[j]++;
+
+                        /* we call model update (in another thread) and then
+                           broadcast the new parameters to member models 
+                           (in another thread) */
+                        if (finishedCount[j] == memberActive->count) {
+                            server->flags[j] = PARAM_STATE_COLLECTED;
+                            if(updater != NULL)
+                                updater->AddJobUpdateSingle(server, memberAll, j, optimizer, broadcaster);
+                        }
+                        else if (finishedCount[j] > memberActive->count) {
+                            ShowNTErrors("Something is wrong with finishedCount!");
+                        }
+                    }
+                }
+            }
+        }
+        else if (collectMode == DATA_COLLECT_REDUCESUM) {
+            for (int j = 0; j < tp.count; j++) {
+
+                bool ready = true;
+
+                /* tp[j]->isGradFinished is true only if the model finishes the computation
+                (in another process) */
+                if (server->flags[j] != PARAM_STATE_NOT_READY || !tp[j]->isGradFinished)
+                    continue;
+
+                /* check if all the models (or part of them) are ready */
+                for (int i = 0; i < memberActive->count; i++) {
+                    XModel * source = (XModel*)memberActive->GetItem(i);
+                    TensorList & sp = source->params;
+
+                    /* sp[j]->isGradFinished is true only if the model finishes the computation
+                    (in another process) */
+                    if (source->flags[j] == PARAM_STATE_COLLECTED ||
+                        source->flags[j] == PARAM_STATE_UPDATED || 
+                        !sp[j]->isGradFinished) 
+                    {
+                        ready = false;
+                        break;
+                    }
+                    else if (source->flags[j] == PARAM_STATE_NOT_READY) {
+                        source->flags[j] = PARAM_STATE_READY;
+                    }
+                }
+
+                if (ready) {
+                    XList tensorList(memberActive->count);
+
+                    for (int i = 0; i < memberActive->count; i++) {
+                        XModel * source = (XModel*)memberActive->GetItem(i);
+                        TensorList & sp = source->params;
+                        tensorList.Add(sp.GetItem(j)->grad);
+                    }
+
+                    /* data transmit */
+                    CollectReduceSum(&tensorList, tp.GetItem(j)->grad);
+
+                    /* reset the flags */
+                    for (int i = 0; i < memberActive->count; i++) {
+                        XModel * source = (XModel*)memberActive->GetItem(i);
+                        source->flags[j] = PARAM_STATE_COLLECTED;
+                    }
+
+                    server->flags[j] = PARAM_STATE_COLLECTED;
+                    finished += memberActive->count;
+
+                    /* we call model update (in another thread) and then
+                       broadcast the new parameters to member models
+                       (in another thread) */
+                    updater->AddJobUpdateSingle(server, memberAll, j, optimizer, broadcaster);
+                }
+            }
+        }
+        else {
+            ShowNTErrors("Unsupported data collection mode!");
+        }
+
+        /* the collection finishes if all data tensors are processed */
+        if (finished == tp.count * memberActive->count)
+            break;
+
+        XSleep(sleepTime);
+    }
+
+    /* reset the flags */
+    //for (int j = 0; j < tp.count; j++)
+    //    server->flags[j] = PARAM_STATE_COLLECTED;
+
+    delete[] finishedCount;
+}
+
+/* 
 collect data 
 >> sourceList - the list of data tensors we collect data from
 >> target - the target tensor we place the result, that is
@@ -68,6 +227,10 @@ void XWorkerCollect::CollectData(XList * sourceList, XModel * target, long sleep
         CheckNTErrors(sp.count == tp.count, "Incompatiable models!");
     }
 
+    /* counts how many member models are collect for each parameters */
+    int * finishedCount = new int[tp.count];
+    memset(finishedCount, 0, sizeof(int) * tp.count);
+
     //fprintf(stderr, "collect data in 0\n");
 
     /* This is a simple implementation of the wait-and-collect process. But
@@ -80,7 +243,7 @@ void XWorkerCollect::CollectData(XList * sourceList, XModel * target, long sleep
             for (int j = 0; j < tp.count; j++) {
                 /* tp[j]->isGradFinished is true only if the model finishes theA computation
                 (in another process) */
-                if (target->flags[j] == PARAM_STATE_COLLECTED || !tp[j]->isGradFinished)
+                if (target->flags[j] != PARAM_STATE_NOT_READY || !tp[j]->isGradFinished)
                     continue;
 
                 /* check if all the models (or part of them) are ready */
@@ -90,7 +253,7 @@ void XWorkerCollect::CollectData(XList * sourceList, XModel * target, long sleep
 
                     /* sp[j]->isGradFinished is true only if the model finishes the computation
                     (in another process) */
-                    if (source->flags[j] != PARAM_STATE_COLLECTED && sp[j]->isGradFinished) {
+                    if (source->flags[j] == PARAM_STATE_NOT_READY && sp[j]->isGradFinished) {
 
                         /* data transmit */
                         CollectP2P(sp[j]->grad, tp[j]->grad);
@@ -98,6 +261,11 @@ void XWorkerCollect::CollectData(XList * sourceList, XModel * target, long sleep
                         /* reset the flag */
                         source->flags[j] = PARAM_STATE_COLLECTED;
                         finished++;
+                        finishedCount[j]++;
+
+                        if (finishedCount[j] == sourceList->count) {
+                            target->flags[j] = PARAM_STATE_COLLECTED;
+                        }
                     }
                 }
             }
@@ -109,7 +277,7 @@ void XWorkerCollect::CollectData(XList * sourceList, XModel * target, long sleep
 
                 /* tp[j]->isGradFinished is true only if the model finishes the computation
                  (in another process) */
-                if (target->flags[j] == PARAM_STATE_COLLECTED || !tp[j]->isGradFinished)
+                if (target->flags[j] != PARAM_STATE_NOT_READY || !tp[j]->isGradFinished)
                     continue;
 
                 /* check if all the models (or part of them) are ready */
@@ -119,7 +287,10 @@ void XWorkerCollect::CollectData(XList * sourceList, XModel * target, long sleep
                     
                     /* sp[j]->isGradFinished is true only if the model finishes the computation
                      (in another process) */
-                    if (source->flags[j] == PARAM_STATE_COLLECTED || !sp[j]->isGradFinished) {
+                    if (source->flags[j] == PARAM_STATE_COLLECTED || 
+                        source->flags[j] == PARAM_STATE_UPDATED || 
+                        !sp[j]->isGradFinished) 
+                    {
                         ready = false;
                         break;
                     }
@@ -163,9 +334,43 @@ void XWorkerCollect::CollectData(XList * sourceList, XModel * target, long sleep
         XSleep(sleepTime);
     }
 
+    //fprintf(stderr, "collect data in 1\n");
+
     /* reset the flags */
-    for (int j = 0; j < tp.count; j++)
-        target->flags[j] = PARAM_STATE_COLLECTED;
+    //for (int j = 0; j < tp.count; j++)
+    //    target->flags[j] = PARAM_STATE_COLLECTED;
+
+    delete[] finishedCount;
+}
+
+/* wrapper of UpdateDataAll */
+void XWorkerCollect::UpdateAll(XList * args)
+{
+    XWorkerCollect * collecter = (XWorkerCollect*)args->GetItem(0);
+    int activeNum = args->GetInt(1);
+    
+    XList memberActive;
+    for (int i = 0; i < activeNum; i++) {
+        XModel * member = (XModel*)args->GetItem(2 + i);
+        memberActive.Add(member);
+    }
+
+    int allNum = args->GetInt(2 + activeNum);
+
+    XList memberAll;
+    for (int i = 0; i < allNum; i++) {
+        XModel * member = (XModel*)args->GetItem(2 + activeNum + 1 + i);
+        memberAll.Add(member);
+    }
+
+    XModel * server = (XModel*)args->GetItem(2 + activeNum + 1 + allNum);
+    XOptimizer * optimizer = (XOptimizer*)args->GetItem(2 + activeNum + 1 + allNum + 1);
+    XWorkerUpdate * updater = (XWorkerUpdate*)args->GetItem(2 + activeNum + 1 + allNum + 2);
+    XWorkerBroadcast * broadcaster = (XWorkerBroadcast*)args->GetItem(2 + activeNum + 1 + allNum + 3);
+
+    collecter->UpdateDataAll(&memberActive, &memberAll, server, 
+                             optimizer, updater, broadcaster, 
+                             SLEEP_TIME_IN_COLLECTING);
 }
 
 /* wrapper of CollectData */
@@ -240,16 +445,58 @@ void XWorkerCollect::CollectAllReduce(XList * all)
 }
 
 /* 
+add a new job of collecting data, update the parameter and 
+broadcast the new parameter
+>> memberActive - member models that are active, i.e., have generated gradients
+>> memberAll -  all member models
+>> server - the server model
+>> optimizer - the optimizer
+>> updater - the worker that updates the parameters
+>> broadcaster - the worker that broadcasts the new parameters to all member
+                 models
+<< return - successful or not
+*/
+bool XWorkerCollect::AddJobUpdateAll(XList * memberActive, XList * memberAll, XModel * server,
+                                     XOptimizer * optimizer, XWorkerUpdate * updater, XWorkerBroadcast * broadcaster)
+{
+    CheckNTErrors(memberActive != NULL, "No input (active) member list!");
+    CheckNTErrors(memberAll != NULL, "No input (all) member list!");
+    CheckNTErrors(server != NULL, "No input server model!");
+    CheckNTErrors(optimizer != NULL, "No input optimizer!");
+    CheckNTErrors(updater != NULL, "No input updater!");
+    CheckNTErrors(broadcaster != NULL, "No input broadcaster!");
+
+    XList args;
+    args.Add(this);
+    args.AddInt(memberActive->count);
+    args.AddList(memberActive);
+    args.AddInt(memberAll->count);
+    args.AddList(memberAll);
+    args.Add(server);
+    args.Add(optimizer);
+    args.Add(updater);
+    args.Add(broadcaster);
+
+    if (isInstantRun)
+        XWorkerCollect::UpdateAll(&args);
+    else
+        queue.EnqueueJob((void*)(char*)XWorkerCollect::UpdateAll, &args);
+
+    return true;
+}
+
+/* 
 add a new job of collecting data
 >> sourceList - the list of models that we want collect data from
 >> target - the destination of the collection
+<< return - successful or not
 */
 bool XWorkerCollect::AddJobCollect(XList * sourceList, XModel * target)
 {
     CheckNTErrors(sourceList != NULL, "no input source model list!");
     CheckNTErrors(target != NULL, "no input target model!");
 
-    XList args;
+    /*XList args;
     args.Add(this);
     args.AddInt(sourceList->count);
     args.AddList(sourceList);
@@ -258,7 +505,22 @@ bool XWorkerCollect::AddJobCollect(XList * sourceList, XModel * target)
     if (isInstantRun)
         XWorkerCollect::Collect(&args);
     else
-        queue.EnqueueJob((void*)(char*)XWorkerCollect::Collect, &args);
+        queue.EnqueueJob((void*)(char*)XWorkerCollect::Collect, &args);*/
+
+    XList args;
+    args.Add(this);
+    args.AddInt(sourceList->count);
+    args.AddList(sourceList);
+    args.AddInt(0);
+    args.Add(target);
+    args.Add(NULL);
+    args.Add(NULL);
+    args.Add(NULL);
+
+    if (isInstantRun)
+        XWorkerCollect::UpdateAll(&args);
+    else
+        queue.EnqueueJob((void*)(char*)XWorkerCollect::UpdateAll, &args);
 
     return true;
 }

source/train/XWorkerCollect.h

@@ -32,6 +32,8 @@
 #include "XWorker.h"
 #include "XModel.h"
 #include "XWorkerJob.h"
+#include "XWorkerUpdate.h"
+#include "XWorkerBroadcast.h"
 
 namespace nts { // namespace nts(NiuTrans.Tensor)
 
@@ -63,9 +65,22 @@ public:
     /* set the collection type */
     void SetCollectMode(DATA_COLLECT_TYPE myMode);
 
+    /* collect the gradient data, update the parameters, and broadcast the 
+       new parameters to all models. NOTE that this method just collect graident
+       from member models. Then it calls an XWorkerUpdate to update the parameters.
+       The XWorkerUpdate also calls an XWorkerBroadcast to broadcast the new parameter
+       to member models back. */
+    void UpdateDataAll(XList * memberActive, XList * memberAll, XModel * server, 
+                       XOptimizer * optimizer, XWorkerUpdate * updater, XWorkerBroadcast * broadcaster, 
+                       long sleepTime);
+
     /* collect the gradient data (i.e., a reducer) */
     void CollectData(XList * sourceList, XModel * target, long sleepTime);
 
+    /* wrapper of UpdateDataAll */
+    static
+    void UpdateAll(XList * args);
+
     /* wrapper of CollectData */
     static
     void Collect(XList * args);
@@ -79,6 +94,10 @@ public:
     /* all-reduce */
     void CollectAllReduce(XList * all);
 
+    /* add a new job of collecting data, update the parameter and broadcast the new parameter */
+    bool AddJobUpdateAll(XList * memberActive, XList * memberAll, XModel * server,
+                         XOptimizer * optimizer, XWorkerUpdate * updater, XWorkerBroadcast * broadcaster);
+
     /* add a new job of collecting data */
     bool AddJobCollect(XList * sourceList, XModel * target);
 

source/train/XWorkerUpdate.cpp

@@ -53,6 +53,37 @@ XOptimizer * XWorkerUpdate::GetOptimizer()
 }
 
 /* 
+update a parameter of a model 
+>> model - the model that we want to update (on the server side)
+>> members - models that would share the updated parameters
+>> pid - the parameter index
+>> optimizer - the optimizer
+>> broadcaster - the worker that would broadcast the new parameter to members
+*/
+void XWorkerUpdate::UpdateParameter(XModel * server, XList * members, int pid,
+                                    XOptimizer * optimizer, XWorkerBroadcast * broadcaster)
+{
+    TensorList & params = server->params;
+    PARAM_STATE * flags = server->flags;
+
+    CheckNTErrors(flags[pid] == PARAM_STATE_COLLECTED, "The state of the parameter is wrong!");
+
+    XTensor * param = params.GetItem(pid);
+    XTensor * grad = param->grad;
+
+    CheckNTErrors(grad != NULL, "No gradient!");
+
+    /* update the parameter */
+    optimizer->UpdateParam(param, grad, pid);
+
+    /* set the flag */
+    flags[pid] = PARAM_STATE_UPDATED;
+
+    /* broadcast the new parameter to other models (in anotehr worker/thread) */
+    broadcaster->AddJobBroadcastSingle(server, members, pid);
+}
+
+/* 
 update the model 
 >> model - the model that we want to update
 >> optimizer - the optimizer
@@ -93,6 +124,31 @@ void XWorkerUpdate::UpdateModel(XModel * model, XOptimizer * optimizer, long sle
 }
 
 /* 
+wrapper of UpdateParameter 
+>> args - arguments of the update
+*/
+void XWorkerUpdate::UpdateSingle(XList * args)
+{
+    CheckNTErrors(args != NULL && args->count >= 6, "Illegal argument list!");
+
+    XWorkerUpdate * updater = (XWorkerUpdate*)args->GetItem(0);
+    XModel * server = (XModel*)args->GetItem(1);
+    int memNum = args->GetInt(2);
+
+    XList members;
+    for (int i = 0; i < memNum; i++) {
+        XModel * member = (XModel*)args->GetItem(3 + i);
+        members.Add(member);
+    }
+
+    int pid = args->GetInt(3 + memNum);
+    XOptimizer * optimizer = (XOptimizer*)args->GetItem(3 + memNum + 1);
+    XWorkerBroadcast * broadcaster = (XWorkerBroadcast*)args->GetItem(3 + memNum + 2);
+
+    updater->UpdateParameter(server, &members, pid, optimizer, broadcaster);
+}
+
+/* 
 wrapper of UpdateModel
 >> args - arguments of the update
 */
@@ -112,6 +168,40 @@ void XWorkerUpdate::Update(XList * args)
 }
 
 /* 
+add a new job of model update (for a parameter) 
+>> model - the model that we want to update (on the server side)
+>> members - models that would share the updated parameters
+>> pid - the parameter index
+>> optimizer - the optimizer
+>> broadcaster - the worker that would broadcast the new parameter to members
+*/
+bool XWorkerUpdate::AddJobUpdateSingle(XModel * model, XList * members, int pid,
+                                       XOptimizer * optimizer, XWorkerBroadcast * broadcaster)
+{
+    CheckNTErrors(model != NULL, "No input model!");
+    CheckNTErrors(members != NULL, "No member model list!");
+    CheckNTErrors(optimizer != NULL, "No optimizer!");
+    CheckNTErrors(broadcaster != NULL, "No broadcaster!");
+    CheckNTErrors(pid >= 0 && pid < model->params.count, "Illegal parameter index!");
+
+    XList args;
+    args.Add(this);
+    args.Add(model);
+    args.AddInt(members->count);
+    args.AddList(members);
+    args.AddInt(pid);
+    args.Add(optimizer);
+    args.Add(broadcaster);
+
+    if (isInstantRun)
+        XWorkerUpdate::UpdateSingle(&args);
+    else
+        queue.EnqueueJob((void*)(char*)XWorkerUpdate::UpdateSingle, &args);
+
+    return true;
+}
+
+/* 
 add a new job of model update
 >> model - the model that we want to update
 >> optimizer - the optimizer

source/train/XWorkerUpdate.h

@@ -30,6 +30,7 @@
 
 #include "XWorker.h"
 #include "XOptimizer.h"
+#include "XWorkerBroadcast.h"
 
 namespace nts { // namespace nts(NiuTrans.Tensor)
 
@@ -55,13 +56,25 @@ public:
     /* get the optimizer */
     XOptimizer * GetOptimizer();
 
+    /* update the parameter */
+    void UpdateParameter(XModel * server, XList * members, int pid,
+                         XOptimizer * optimizer, XWorkerBroadcast * broadcaster);
+
     /* update the model */
     void UpdateModel(XModel * model, XOptimizer * optimizer, long sleepTime);
 
+    /* wrapper of UpdateParameter */
+    static
+    void UpdateSingle(XList * args);
+
     /* wrapper of UpdateModel */
     static
     void Update(XList * args);
 
+    /* add a new job of model update (for a parameter) */
+    bool AddJobUpdateSingle(XModel * model, XList * members, int pid,
+                            XOptimizer * optimizer, XWorkerBroadcast * broadcaster);
+
     /* add a new job of model update */
     bool AddJobUpdate(XModel * model, XOptimizer * optimizer);
 };