Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Train in half precission #257

Closed
hvico opened this issue Aug 26, 2021 · 2 comments
Closed

Train in half precission #257

hvico opened this issue Aug 26, 2021 · 2 comments
Labels
enhancement New feature or request

Comments

@hvico
Copy link

hvico commented Aug 26, 2021

Hello,

It would be great to have a training mode or arg which would enable FP16 mode, for enhanced performance on CUDA traininig.

I use this training code for regular hugginface training:

`scaler = torch.cuda.amp.GradScaler()

def train_fp16(epoch):
model.train()
for _,data in enumerate(training_loader, 0):
ids = data['ids'].to(device, dtype = torch.long)
mask = data['mask'].to(device, dtype = torch.long)
token_type_ids = data['token_type_ids'].to(device, dtype = torch.long)
targets = data['targets'].to(device, dtype = torch.float)

    with torch.cuda.amp.autocast():

        outputs = model(ids, mask, token_type_ids)

    optimizer.zero_grad()
    loss = loss_fn(outputs, targets)
    #if _%1000==0:
    #    print(f'Epoch: {epoch}, Loss:  {loss.item()}')
    
    optimizer.zero_grad()
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

def train_fp32(epoch):
model.train()
for _,data in enumerate(training_loader, 0):
ids = data['ids'].to(device, dtype = torch.long)
mask = data['mask'].to(device, dtype = torch.long)
token_type_ids = data['token_type_ids'].to(device, dtype = torch.long)
targets = data['targets'].to(device, dtype = torch.float)

    outputs = model(ids, mask, token_type_ids)

    optimizer.zero_grad()
    loss = loss_fn(outputs, targets)
    if _%1000==0:
        print(f'Epoch: {epoch}, Loss:  {loss.item()}')
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()`

Thanks!

@EricFillion EricFillion added the enhancement New feature or request label Sep 11, 2021
@EricFillion
Copy link
Owner

Thanks for the suggestion! I'll look into it. I agree with adding more features to reduce memory consumption.

@EricFillion
Copy link
Owner

Once again, thanks for this suggestion. I just published a new version that allows you to enable fp16.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

No branches or pull requests

2 participants