Usage of Pytorch Native AMP in place of apex (Pytorch 1.6) in Trainer

[prajjwal1]

🚀 Feature request

It would be nice to remove the Apex dependency for fp16 training and use the native Pytorch's AMP methods in the Trainer method. Pytorch recommends Apex users to switch to it's native implementation, even Apex does it so. Moreover, it will eliminate the need for users to build apex by themselves.

Your contribution

I am happy to submit a PR if you think it would be a good addition. Please let me know.

[prajjwal1]

I had a query, Pytorch examples show loss being calculated as :

with autocast():
    output = model(input)
    loss = loss_fn(output, target)
scaler.scale(loss).backward()

But in all SequenceClassification and other models, loss is calculated in the forward pass. We can use @autocast decorator on the forward pass as the docs suggest, but this introduce so many changes for one feature. Maybe, there's a workaround. Does computing loss in autocast scope affect the loss itself when backward is called upon it ?

[sgugger]

Hi there,

Note that we won't pin the version of PyTorch to 1.6 minimum, so the use of native mixed precision will have to be controlled by a test on the pytorch version (basically use native mixed precision when the version allows it and use apex otherwise).

Otherwise, I don't think the loss being computed inside the model should be a problem, the line would probably be

with autocast():
    outputs = model(**inputs)
    loss = outputs[0]

inside Trainer but I haven't run tests yet.

You're welcome to try to work on a PR with this, otherwise it is on my TODO for when I have time (hopefully in the next few weeks).

[prajjwal1]

Hi Sylvain,

I've opened up a PR. I know that pinning of version won't be done. To address, I've addressed this issue the same way as we handle scheduler.get_lr().