Xla: Training on multiple core performs badly while training on one xla core is good

Created on 6 Aug 2019  Â·  23Comments  Â·  Source: pytorch/xla

the code is as follows:

devices = xm.get_xla_supported_devices(max_devices=N)   # N is set to 1 or 8

model_parallel = dp.DataParallel(model, device_ids=devices)
def train_loop_fn(model, loader, device, context):
     optimizer = context.getattr_or(
               'optimizer', 
                torch.optim.Adam(
                    filter(lambda param: param.requires_grad, model.parameters()), 
                    lr=p["lr"] * len(devices), ))
     model.train()
     for i, (inp, targ) in loader:
          optimizer.zero_grad()
          outp = model(inp)
          loss = loss_fn(targ, outp)
          loss.backward()
          xm.optimizer_step(optimizer)

model_parallel (train_leep_fn, train_loader)

p.s. I see in the test_train_cifar and test_train_imagenet.py example, lr is not scaled to w.r.t number of cores while in test_train_mnist.py it model_parallel . wondering which way should we obey?

also, not sure if I understand the code correclty, from the data_parallel.py it seems a different copy of the model is placed on each device with their own optimizer and get updated _independently_. if this's the case, then wondering how the training benefits from the multi-core? apologize if I missed anything

Most helpful comment

image

All 23 comments

The gradients are synced before updating, so they do get updated the same (this happens in xm.optimizer_step). In the end, we have n models w/ the same exact weights.

Like Taylan mentioned, gradients should sync. I once had a bug where gradients did not sync, but I was using a mismatch of pytorch version and TPU version, e.g. pytorch-r0.1 conda env with a pytorch-nightly TPU or vice-versa

thanks for the reply! then let me recheck the code maybe.

How about the learning rate then? is it supposed to be scaled or not?

@zcain117 wondering how to check the version of TPU?

@zcain117 wondering how to check the version of TPU?

You can check either on the Google Cloud TPU UI or via the CLI with something like this. Check for the TPU Software version.

image

About the difference in performance between multiple cores vs. one xla core. By performance, I am assuming you are referring to how quickly the loss drops after each epoch. I have noticed the same problem with mnist and other models with GPUs and with TPUs. So it seems this is a general problem for multi-core training, not an xla problem. I think this only happens when the training data is very small (as with mnist).

About the difference in performance between multiple cores vs. one xla core. By performance, I am assuming you are referring to how quickly the loss drops after each epoch. I have noticed the same problem with mnist and other models with GPUs and with TPUs. So it seems this is a general problem for multi-core training, not an xla problem. I think this only happens when the training data is very small (as with mnist).

thanks for that suggestion, maybe i'll let it run longer and see if it's getting better xd

oh I see the TPU version, it's pytorch-0.1, so does that mean I cannot run the project under pytorch-nightly environment?

if you have the same batch size;

in 1 epoch, with 1 core, you will do N steps
whereas in 1 epoch, with 8 cores, you will do N/8 steps.

Effectively, 1 step in 8 core training is; computing 8 gradients and
syncing them, so it is analogous to training on 1 core with 8x the batch
size.

On Tue, Aug 6, 2019, 17:33 Iz Beltagy notifications@github.com wrote:

About the difference in performance between multiple cores vs. one xla
core. By performance, I am assuming you are referring to how quickly the
loss drops after each epoch. I have noticed the same problem with mnist and
other models with GPUs and with TPUs. So it seems this is a general problem
for multi-core training, not an xla problem. I think this only happens when
the training data is very small (as with mnist).

—
You are receiving this because you commented.
Reply to this email directly, view it on GitHub
https://github.com/pytorch/xla/issues/907?email_source=notifications&email_token=ABDAQYLBV2LIHEUCL2OQOELQDIJ4ZA5CNFSM4IJYGSCKYY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOD3W23DY#issuecomment-518892943,
or mute the thread
https://github.com/notifications/unsubscribe-auth/ABDAQYK7RQ45EDNYV7MT5MTQDIJ4ZANCNFSM4IJYGSCA
.

@taylanbil thanks i see how the mult-core works together now. yet how is this related to the learning rate's scale?

also, if we r regarding it as training 8xbatch each step, is it that I'm supposed to set the batchsize in multicore training to the 1/8 of the original batch size?

learning rate for all optimizers needs to change in sync as well. Now I
have not verified if this already happens myself, but I can get back to you
on that tomorrow.

as a rule of thumb, tpus work well with larger batch size. In the end, it
is a parameter you need to tune in your training job.

There is no need to do 1/8 the batch size generally. I was just pointing
out the comparison between single and multi core is "unfair" if you dont do
that.

Extreme case of this is when batch size=dataset size/8. with 1 core, you
will take 8 steps, and with 8 cores, you will take just 1 step. So it is
super likely that after epoch 1, singlecore will have smaller loss.

On Tue, Aug 6, 2019, 19:58 Crystina notifications@github.com wrote:

@taylanbil https://github.com/taylanbil thanks i see how the mult-core
works together now. yet how is this related to the learning rate's scale?

also, if we r regarding it as training 8xbatch each step, is it that I'm
supposed to set the batchsize in multicore training to the 1/8 of the
original batch size?

—
You are receiving this because you were mentioned.
Reply to this email directly, view it on GitHub
https://github.com/pytorch/xla/issues/907?email_source=notifications&email_token=ABDAQYLPQ2OT62UU4FVUUK3QDI26PA5CNFSM4IJYGSCKYY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOD3XBMLQ#issuecomment-518919726,
or mute the thread
https://github.com/notifications/unsubscribe-auth/ABDAQYPULHHO3IG4AUFQQKDQDI26PANCNFSM4IJYGSCA
.

oh i see now, that's a bit different with what I though the multi-core is doing: I thought each batch of data is split on each tpu core, then the losses r aggregated to do the backpropagation. nvm i'm using different batch size to train the multi-core model. see if this solve the problem. Thanks a lot for the elaboration.

and by the scale of learning rate I mean, in mnist example, it multiplies the lr with the number of cores (here) yet it's not doing so in imagenet or other examples. so I'm not sure which is the correct way to do. (in imagenet)

re: learning rates in multi core: In my example, learning rates seem to be synced.

In PyTorch/TPU the batch size is the per-core batch size. We sometime call it mini batch size internally.

Learning rate scaling WRT batch size.
Usually a bigger batch allows larger learning rates, as your gradients gather signal from more samples.

oh I see the TPU version, it's pytorch-0.1, so does that mean I cannot run the project under pytorch-nightly environment?

It might be fine or you might run into weird issues in my experience. It's safer to use the same version for TPU that you use for the environment. You could bring down the current TPU and make a nightly TPU instead--see that UI screenshot that @ibeltagy showed above for the dropdown menu to choose pytorch version during TPU creation

thanks so much for all the suggestions and explanation above, i'll go back to pytorch-0.1 version with scaled lr, and try difference batch sizes for multicore. get back the result to this channel later!

oh I guess the problem is solved: partly due to the batch size partly due to the evaluation metric implementation - I didn't integrate the eval result from each xla device well. Now the results on multi-core is consistent with single core!

and just for reference, the result on pytorch-night and pytorch-0.1 are consistent.

Thanks a lot for all the help. closing this issue now xd

@taylanbil I understand that for Pytorch/GPU, batch size is per-core batch size but I'm still not sure abt learning rate. So, if for a single core my learning rate is x, then for 8 cores, should it be x*8?

Not always the case, but increasing the global batch size increases the confidence you can assume in moving with (inverse) gradient step.
IOW you can afford taking longer steps.
Whether 8x or lower, depends on the model and only hyper-parameters tuning can tell you.

Assuming your gradient is additive, @ajay960singh if you wanted to exactly reproduce "the result you get with 1 core" on 8 cores, then:

  • your learning rate should be the same
  • your batch size on 8 cores should be 1/8th of the batch size in 1 core

However, this is not practically a good idea. The beauty of having multiple cores at your disposal is, you _can_ afford bigger batches per step (keeping batch size constant per core, 1 host can process 8x bigger batches in one step since it has 8 cores).

And as Davide points out above, if you process more samples in one step, that step's gradient has more information and less noise, so taking a bigger step (via a larger learning rate) is a good idea.

In the end it's a tunable parameter, and it depends on the model/task/dataset whether or not 8x is best or 2x is best or 20x is best.

Thanks for the elaborate explanation guys. That was of huge help!

Thanks for this issue! Lot my doubts/misunderstandings were resolved :) Thanks guys!

Was this page helpful?
0 / 5 - 0 ratings

Related issues

patrickvonplaten picture patrickvonplaten  Â·  7Comments

myleott picture myleott  Â·  7Comments

myleott picture myleott  Â·  6Comments

Arjuna197 picture Arjuna197  Â·  5Comments

ibeltagy picture ibeltagy  Â·  5Comments