Design Doc: Synchronous SGD

  1. # Basic design philosophy:
  3. #
  5. # - Workers are gRPC client, the master is the gRPC server.
  7. # - The master keeps the three task queues: todo, doing, and done.
  9. # - The master keeps the current global model and the model ID is the hash of the model parameters.
  11. #------- worker.py -------#
  13. model_params = NULL
  14. model_version = NULL
  16. module = parse_command_line_for_class_name()(
  17.     parse_command_line_for_ctor_params())
  19. master = grpc.create_client(getenv("MASTER_ADDR"))
  21. while True:
  22.     # Claim a task from the master. A task consists of a data segment and the
  23.     # model id.  The meaning of a task is "to update the specified model with
  24.     # the given data segment".
  25.     task, err = master.GetTask()
  26.     if err == NO_MORE_TASK:
  27.         break    # Training completed.
  28.     if err != NULL:
  29.         continue # Retry to get task.
  31.     task_status = SUCCEED
  32.     for minibatch in read_data(task):
  33.         accepted = False
  34.         report_count = 0
  35.         while not accepted:
  36.             try:
  37.                 # If the current model_version on the worker is older than the model
  38.                 # on the master, this call updates model_version and
  39.                 # model_params; otherwise, it leaves these two variables unchanged.
  40.                 master.UpdateModelIfOutOfDate(&model_version, &model_params)
  41.                 cost = module.forward(data, model_params)
  42.                 gradients = module.backward(cost, model_params)
  43.             except:
  44.                 task_status = FAILED
  45.                 break
  46.             else:
  47.                 # If the reported gradients are not accepted by the master due to old model_version,
  48.                 # try the minibatch again with the updated model in the next while loop.
  49.                 # Fail the task if the minibatch report count exceeds a predefined threshold.
  50.                 accepted = master.ReportGradients(model_version, gradients)
  51.                 if not accepted:
  52.                     report_count += 1
  53.                     if report_count == PREDEFINED_MAX_REPORT_COUNT:
  54.                         task_status = FAILED
  55.                         break
  56.         if task_status == FAILED:
  57.             break
  58.     master.ReportTask(task, task_status)
  60. #------- master.py -------#
  62. inputs = parse_command_line().recordio_files() # No online learning in this version.
  63. todo = partition_data_into_tasks(inputs)
  64. doing = Queue()
  65. done = Queue()
  67. module = parse_command_line_for_class_name()(
  68.     parse_command_line_for_ctor_params())
  70. model_params = module.create_parameters().random_initialize()
  71. model_version = 0
  73. gradients = []
  75. @grpc
  76. def UpdateModelIfOutOfDate(mv, mp):
  77.     if model_version != mv:
  78.         copy(*mv, model_version)
  79.         copy(*mp, model_params)
  81. @grpc
  82. def GetTask():
  83.     task = todo.pop()
  84.     doing.push(task)
  85.     return task
  87. @grpc
  88. def ReportGradients(mv, grads):
  89.     accepted = False
  90.     if mv == model_version:
  91.         gradients += grads
  92.         accepted = True
  93.         if len(gradients) >= num_gradients_sufficient_to_update_model():
  94.             model_params = optimize_model(model_params, gradients)
  95.             model_version = model_version + 1
  96.             gradients = [] # Clear out the buffer.
  97.     return accepted
  99. @grpc
  100. def ReportTask(task, status):
  101.     if status == FAILED:
  102.         move_task(task, doing, todo) # Move the task from doing back to todo
  103.     else:
  104.         move_task(task, doing, done) # Move the task from doing to done