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,
 bool XLeader::Run(XConfig * config, DataDistributeBase * dataDistributor,
                   XModel * model, XOptimizer * optimizer)
                   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;
     bool isDataOK = true;
     int activeJobCount = 0;
     int activeJobCount = 0;
     int* active = new int[jworkers.count];
     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 */
         /* member models that are active in this run */
         XList members(jworkers.count);
         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
         /* jobs in queue 2: collect the (gradient) data and other stuff. This
            is a reduce process. */
            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 */
         /* 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
         /* 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
            we would update a worker to the laster model parameters, even if it is
            not involved in this run. */
            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();
         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
 >> source - the data that we want to broadcast
 >> targetList - the target places that we recieve the data
 >> targetList - the target places that we recieve the data
 >> sleepTime - the waiting time in broadcasting
 >> sleepTime - the waiting time in broadcasting
@@ -72,17 +97,12 @@ void XWorkerBroadcast::BroadcastData(XModel * source, XList * targetList, long s
     while (1) {
     while (1) {
         for (int i = 0; i < sp.count; i++) {
         for (int i = 0; i < sp.count; i++) {
             if (source->flags[i] == PARAM_STATE_UPDATED && finishedFlag[i] == 0) {
             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;
                 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 
 wrapper of BroadcastData 
 >> args - the list of arguments
 >> 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(target != NULL, "The target tensor should not be NULL!");
     CheckNTErrors(IsSameShaped(*source, *target), "The two tensors should be of the same shape!");
     CheckNTErrors(IsSameShaped(*source, *target), "The two tensors should be of the same shape!");
 
 
+    if(source != target)
         CopyValues(*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
 >> source - the data that we want to broadcast
 >> targetList - the target places that we recieve the data
 >> targetList - the target places that we recieve the data
 */
 */

source/train/XWorkerBroadcast.h

@@ -60,9 +60,16 @@ public:
     /* set the broadcasting type */
     /* set the broadcasting type */
     void SetBroadcastMode(DATA_BROADCAST_TYPE myMode);
     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);
     void BroadcastData(XModel * source, XList * targetList, long sleepTime);
 
 
+    /* wrapper of BroadcastDataSingle */
+    static
+    void BroadcastSingle(XList * args);
+
     /* wrapper of BroadcastData */
     /* wrapper of BroadcastData */
     static
     static
     void Broadcast(XList * args);
     void Broadcast(XList * args);
@@ -70,7 +77,10 @@ public:
     /* P2P data broadcasting */
     /* P2P data broadcasting */
     void BroadcastP2P(XTensor * source, XTensor * target);
     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);
     bool AddJobBroadcast(XModel * source, XList * targetList);
 };
 };
 
 

source/train/XWorkerCollect.h

@@ -32,6 +32,8 @@
 #include "XWorker.h"
 #include "XWorker.h"
 #include "XModel.h"
 #include "XModel.h"
 #include "XWorkerJob.h"
 #include "XWorkerJob.h"
+#include "XWorkerUpdate.h"
+#include "XWorkerBroadcast.h"
 
 
 namespace nts { // namespace nts(NiuTrans.Tensor)
 namespace nts { // namespace nts(NiuTrans.Tensor)
 
 
@@ -63,9 +65,22 @@ public:
     /* set the collection type */
     /* set the collection type */
     void SetCollectMode(DATA_COLLECT_TYPE myMode);
     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) */
     /* collect the gradient data (i.e., a reducer) */
     void CollectData(XList * sourceList, XModel * target, long sleepTime);
     void CollectData(XList * sourceList, XModel * target, long sleepTime);
 
 
+    /* wrapper of UpdateDataAll */
+    static
+    void UpdateAll(XList * args);
+
     /* wrapper of CollectData */
     /* wrapper of CollectData */
     static
     static
     void Collect(XList * args);
     void Collect(XList * args);
@@ -79,6 +94,10 @@ public:
     /* all-reduce */
     /* all-reduce */
     void CollectAllReduce(XList * all);
     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 */
     /* add a new job of collecting data */
     bool AddJobCollect(XList * sourceList, XModel * target);
     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 
 update the model 
 >> model - the model that we want to update
 >> model - the model that we want to update
 >> optimizer - the optimizer
 >> 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
 wrapper of UpdateModel
 >> args - arguments of the update
 >> 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
 add a new job of model update
 >> model - the model that we want to update
 >> model - the model that we want to update
 >> optimizer - the optimizer
 >> optimizer - the optimizer

source/train/XWorkerUpdate.h

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