Commit df6e3cd7 authored by mshoeybi's avatar mshoeybi
Browse files

set grads to none for the contig. buffers

parent 68797d90

megatron/training.py

+1 −2
Original line number Diff line number Diff line
@@ -354,7 +354,6 @@ def train_step(forward_step_func, data_iterator, 
    if args.DDP_impl == 'local' and args.use_contiguous_buffers_in_ddp:

        for partition in model:

            partition.zero_grad_buffer()

    else:

    optimizer.zero_grad()



    forward_backward_func = get_forward_backward_func()