tensorflow: Problem about distributed training with XLA compiling.

‘jit_compile’ is the new alias for ‘experimental_compile’

Current suggestion is to “jit_compile” only parts of training which are running independently on each replica(GPU). Any time any communication/synchronization is needed then jit_compile around that will fail. So a ‘jit_compile’ around entire strategy.run will fail, and a ‘jit_compile’ around the function containing ‘optimizer.apply_gradients’ will fail. Also any ‘jit_compile’ around functions which update metrics will fail. But ‘jit_compile’ around the function doing the main training should work.

We are looking to lower the all_reduce and resolved some of these issues. Soon folks can enable ‘jit_compile’ on the entire training step (ideally like the way the user has done in this case).

To fix their current code, user needs to change the train_step code in <their github>/main/op_utils.py from:

    @tf.function(experimental_compile = args.compile)
    def train_step(images, labels):
        with tf.GradientTape() as tape:
            pred = model(images, training = True)
            total_loss = loss_object(labels, pred)/args.batch_size
        gradients = tape.gradient(total_loss, model.trainable_variables)
        if args.weight_decay > 0.:
            gradients = [g+v*args.weight_decay for g,v in zip(gradients, model.trainable_variables)]
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))

        train_loss.update_state(total_loss)
        train_accuracy.update_state(labels, pred)
   
    @tf.function(experimental_compile = args.compile)
    def train_step_dist(image, labels):
        strategy.run(train_step, args= (image, labels))

Code needs to changed to:

    @tf.function(jit_compile = True)
    def compiled_step(images, labels):
        with tf.GradientTape() as tape:
            pred = model(images, training = True)
            total_loss = loss_object(labels, pred)/args.batch_size
        gradients = tape.gradient(total_loss, model.trainable_variables)
        return total_loss, pred, gradients

    def train_step(images, labels):
        total_loss, pred, gradients = compiled_step(images, labels)
        if args.weight_decay > 0.:
            gradients = [g+v*args.weight_decay for g,v in zip(gradients, model.trainable_variables)]

        optimizer.apply_gradients(zip(gradients, model.trainable_variables))

        train_loss.update_state(total_loss)
        train_accuracy.update_state(labels, pred)
   
    @tf.function()
    def train_step_dist(image, labels):
        strategy.run(train_step, args= (image, labels))

Please let me know if you still hit issues, and we will be happy to resolve them.

For runs on single-host (single or multi-GPU using mirrored-strategy) we now support XLA around optimizer.apply_gradients . So now for even better performance we can do:

    @tf.function(jit_compile=True)
    def train_step(*data):
        # no need for xla around compiled_step 
        total_loss, gradients, update_vars = compiled_step(*data)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))

+@cheshire

@sseung0703 Your example does not show the whole code, since the part actually applying experimental_compile=True is not in your example.

@nnigania is currently working on supporting collective ops under XLA:GPU compilation, so I think the best bet is to wait for this work to land.

Alternatively, you would have to compile the part of your model which is not using any collectives (as we did for MLPerf).