search

DifferentiableUniverseInitiative / mesh

Mesh TensorFlow: Model Parallelism Made Easier
Apache License 2.0
0 stars 0 forks source link

Adding support for hvd LaidOutVariable #1

Open EiffL opened 3 years ago

EiffL commented 3 years ago

In our current prototype of the horovod mesh implementation, we are using copy/pasted code from the TPU SIMD implementation, it is not expected to work... https://github.com/DifferentiableUniverseInitiative/mesh/blob/70a13d38c5b4b16200dcb8f3d68f866633875181/mesh_tensorflow/hvd_simd_mesh_impl.py#L112

This issue is to document the reimplementation of these variables using the horovod backend. As points of reference, we can look at these variables are implemented in both the Device Placemnt Impl and SIMD impl

EiffL commented 3 years ago

The goal here would be to be able to train the MNIST demo model with the new backend implemenation https://github.com/DifferentiableUniverseInitiative/mesh/blob/master/examples/mnist.py

tobias-liaudat commented 3 years ago

Hi @EiffL, I'm leaving this comment as a checkpoint of when we finished the IDRIS hackaton.

With @b-remy we were trying to solve this issue but we hadn't had the time to finish. I uploaded my progress in the branch tob_vars and Benjamin on his branch ben-variable.

With our current implementation we were able to run a forward pass of a simple dense network. However, when we tried to optimise the network we were having an error and we're not sure of the origin and how to solve it.

The error we found was this one, and is always related with the Assign and the optimisation.

Traceback (most recent call last):
  File "optim_demo.py", line 174, in <module>
    (type(obj).__name__, types_str))
TypeError: Can not convert a Assign into a Tensor or Operation.
TypeError: Fetch argument <mesh_tensorflow.ops.Assign object at 0x14d1fcb94c50> has invalid type <class 'mesh_tensorflow.ops.Assign'>, must be a string or Tensor. (Can not convert a Assign into a Tensor or Operation.)

To reproduce the error one can run this test script with this job script. The implementation of the LaidOutVariable is here.

EiffL commented 3 years ago

Thanks so much @tobias-liaudat will take a look!

EiffL commented 3 years ago

ok found a couple of problems, the update seems to work now, will create a proper branch ^^

EiffL commented 3 years ago

I've opened branch variables starting from tob_vars, and cleaned it up a bit. Problems where:

EiffL commented 3 years ago

probably, we can try to first train the mnist model with commenting out the restore and save parts

EiffL commented 3 years ago

oookkkkk so I tried something here: https://github.com/DifferentiableUniverseInitiative/mesh/tree/u/EiffL/toy_model

with this script https://github.com/DifferentiableUniverseInitiative/mesh/blob/u/EiffL/toy_model/examples/toy_model_gpu.sh

It runs apparently, can save and restore, but not clear if it's actually training ^^" the loss function doesnt go down much