Skip to content

GitLab

Commit 28cd66e1 authored by mohammad's avatar mohammad Committed by Deepak Narayanan
Browse files

fp32 working

parent b4b0d739
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
megatron/training.py
View file @ 28cd66e1
...@@ -238,10 +238,11 @@ def backward_step(optimizer, model, loss): ...@@ -238,10 +238,11 @@ def backward_step(optimizer, model, loss):
# Backward pass. # Backward pass.
timers('backward-backward').start() timers('backward-backward').start()
optimizer.zero_grad(set_grads_to_None=True)
if args.fp16: if args.fp16:
optimizer.zero_grad(set_grads_to_None=True)
optimizer.backward(loss, update_master_grads=False) optimizer.backward(loss, update_master_grads=False)
else: else:
optimizer.zero_grad()
loss.backward() loss.backward()
timers('backward-backward').stop() timers('backward-backward').stop()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment