Xla: Cannot train models with >1B parameters

[myleott]

🐛 Bug

When trying to train a model with >1B parameters I get an error:
assertion failed: byte_count_ < total_size_

This issue suggests it may be hitting a protobuf limit of 2GB. I suspect there's some proto somewhere corresponding to the bfloat16 params, which goes over this 2GB limit.

Any ideas or possible workarounds?

To Reproduce

# 1.0B params, success:
python repro.py --num-embed 50000 --embed-dim 1536 --num-layers 24
# 1.2B params, assertion failed: byte_count_ < total_size_:
python repro.py --num-embed 50000 --embed-dim 1536 --num-layers 30

# repro.py
import argparse

import torch
import torch.nn as nn

import torch_xla.core.xla_model as xm
import torch_xla.distributed.parallel_loader as pl
import torch_xla.distributed.xla_multiprocessing as xmp

class Net(nn.Module):

    def __init__(self, num_embed=50000, embed_dim=1024, num_layers=24):
        super().__init__()
        self.embed = nn.Embedding(
            num_embeddings=num_embed, embedding_dim=embed_dim, padding_idx=0
        )
        self.layers_a = nn.ModuleList([
            nn.Sequential(
                nn.LayerNorm(embed_dim),
                nn.Linear(embed_dim, 3*embed_dim),  # q, k, v input projection
                nn.Linear(3*embed_dim, embed_dim),  # skip self-attention
                nn.Linear(embed_dim, embed_dim),    # output projection
                nn.Dropout(),
            )
            for i in range(num_layers)
        ])
        self.layers_b = nn.ModuleList([
            nn.Sequential(
                nn.LayerNorm(embed_dim),
                nn.Linear(embed_dim, 4*embed_dim),  # FFN
                nn.ReLU(),
                nn.Linear(4*embed_dim, embed_dim),  # FFN
                nn.Dropout(0.1),
            )
            for i in range(num_layers)
        ])
        self.out_proj = nn.Linear(embed_dim, num_embed)

    def forward(self, tokens):
        x = self.embed(tokens)
        for layer_a, layer_b in zip(self.layers_a, self.layers_b):
            x = x + layer_a(x)
            x = x + layer_b(x)
        x = self.out_proj(x)
        return x

class DataGenerator():
    def __init__(self, length=1024, bsz=1):
        self.seq = torch.arange(length).repeat(bsz, 1)
    def __iter__(self):
        return self
    def __next__(self):
        return self.seq

def main(rank):
    parser = argparse.ArgumentParser()
    parser.add_argument('--num-embed', type=int)
    parser.add_argument('--embed-dim', type=int)
    parser.add_argument('--num-layers', type=int)
    args = parser.parse_args()

    model = Net(args.num_embed, args.embed_dim, args.num_layers)
    num_params = sum(p.numel() for p in model.parameters())
    print('num_params: {}'.format(num_params))

    device = xm.xla_device()
    model = model.to(device=device, dtype=torch.bfloat16)
    tpu_loader = pl.ParallelLoader(DataGenerator(), [device]).per_device_loader(device)

    next(tpu_loader)  # fails for large models

    print('success')


if __name__ == "__main__":
    xmp.spawn(main, args=(), nprocs=1)

Environment

• reproducible on XLA backend [CPU/TPU]: TPU
• torch_xla version: master

[dlibenzi]

That is a sad protobuf story, but yes, the limit is there.
Can you split into different tensors?

[myleott]

Actually it seems to be an issue with the ParallelLoader, which I guess calls xm.mark_step internally.

I've updated the repro to a model with 1.2B params, which are no longer all in a single tensor, and it fails on next(tpu_loader).

[dlibenzi]

I think I know what is going on 😄
The ParallelLoader tries to batch tensor upload for efficiency, and the whole batch ultimately ends up within a single ProtoBuf when hitting the TPU VM.
Let me fix that.

[myleott]

Sure, you know best :) Although it's surprising that the smaller model (1B params) works with the same batch input, while the slightly larger model (1.2B params) doesn't.

[taylanbil]

The initial repro script crashes during forward pass, w/o parallelloader or anything. These two crashes (first repro vs second repro) are different imho.

# repro.py
import argparse

import torch
import torch.nn as nn
import torch.optim as optim

import torch_xla.core.xla_model as xm
import torch_xla.distributed.xla_multiprocessing as xmp

class Net(nn.Module):
    def __init__(self, nparams):
        super().__init__()
        self.layer = nn.Parameter(torch.zeros(nparams))

    def forward(self):
        return self.layer.sum()

def main(rank):
    parser = argparse.ArgumentParser()
    parser.add_argument('--nparams', type=int)
    args = parser.parse_args()

    device = xm.xla_device()
    model = Net(args.nparams).to(device=device, dtype=torch.bfloat16)
    optimizer = optim.SGD(model.parameters(), lr=0.001)
    loss = model()  # crashes here.
    loss.backward()
    xm.optimizer_step(optimizer)

if __name__ == "__main__":
    main(0)

[myleott]

These two crashes (first repro vs second repro) are different imho.

Yes, you're right. The current one (second repro) is the real problem affecting me. The original one (first repro) is probably a less common issue we can ignore for now.

Note that the original one (first repro) is solved by splitting the Parameter into two Parameters each with 0.5B values:

self.layer1 = nn.Parameter(torch.zeros(nparams // 2))
self.layer2 = nn.Parameter(torch.zeros(nparams // 2))

[taylanbil]

This fixes it too as Davide suggested, while leaving model structure intact.

# repro.py
import argparse

import torch
import torch.nn as nn
import torch.optim as optim

import torch_xla.core.xla_model as xm
import torch_xla.distributed.xla_multiprocessing as xmp

class Net(nn.Module):
    def __init__(self, nparams):
        super().__init__()
        self.layer = nn.Parameter(torch.zeros(nparams))

    def forward(self):
        return self.layer.sum()


def splitto(model, args, device):
    origshape = model.layer.shape
    splittensors = model.layer.split(args.nparams // 2)
    splittensors = [
        st.to(device=device, dtype=torch.bfloat16) for st in splittensors
    ]
    model.layer = nn.Parameter(torch.stack(splittensors).reshape(origshape))


def main(rank):
    parser = argparse.ArgumentParser()
    parser.add_argument('--nparams', type=int)
    args = parser.parse_args()
    device = xm.xla_device()
    model = Net(args.nparams)
    splitto(model, args, device)
    #model = model.to(device=device, dtype=torch.bfloat16)
    optimizer = optim.SGD(model.parameters(), lr=0.001)
    loss = model()
    loss.backward()
    xm.optimizer_step(optimizer)

if __name__ == "__main__":
    main(0)

[dlibenzi]

@taylanbil Can you try this in your repro:

xm.optimizer_step(optimizer, barrier=True)

And:

export XLA_SYNC_WAIT=1

[taylanbil]

added barrier and ran

$ TF_CPP_VMODULE=tensor=5 XLA_SYNC_WAIT=1 python repro.py --nparams 1100000000

2020-02-13 20:51:30.808727: I torch_xla/csrc/tensor.cpp:1235] 3 live tensors: devices=[]
2020-02-13 20:51:30.808951: I torch_xla/csrc/tensor.cpp:987] Waiting on device barrier for device TPU:0 ...
2020-02-13 20:51:30.808984: I torch_xla/csrc/tensor.cpp:993] Waiting on device barrier for device TPU:0 done!
2020-02-13 20:51:30.809038: I torch_xla/csrc/tensor.cpp:1031] Tensors graph hash 11667419629977000337 on device TPU:0
2020-02-13 20:51:30.809321: I torch_xla/csrc/tensor.cpp:1342] Compiling IR graph hash 11667419629977000337 on device TPU:0 ...  # stuck here at the moment.

It takes a while to compile w/ XLA_SYNC_WAIT

update: It's been ~40 mins, still stuck.

[taylanbil]

Actually, even w/o XLA_SYNC_WAIT, adding barrier=True makes it slow.

[taylanbil]

So, the takeaway from the first repro here is, that XLA does not work well with single layers with this many parameters. But it sounds like this is not the actual issue in practice anyway, and that's being handled by Davide.

[dlibenzi]

https://github.com/pytorch/xla/pull/1643

This could be fixed even by reducing the ParallelLoader device_prefetch_size argument from 4 down to the necessary value.

[dlibenzi]

There is another fix into the C++ side, we have to do.
We will fix tomorrow ...

[dlibenzi]

https://github.com/pytorch/xla/pull/1646 should fix that issue.
@taylanbil will be checking with that script directly.
I tested locally setting XRT_MAX_TENSORS_PARTITION to small value and the XrtPartitionedTransferToServer is ticking.

[taylanbil]

Tested w/ repro script from above, works !

[shz0116]

It seems the 2GB limit is still there. I use EmbeddingBag(30000000, 128), I got the same problem.
Any suggestion ?

[davidel]

We got around that in cases where a single tensor is not >2GB.
If you have a single tensor >2GB, can't be fixed as the limit is deep down into the Protocol Buffers core API.

[shz0116]

We got around that in cases where a single tensor is not >2GB.
If you have a single tensor >2GB, can't be fixed as the limit is deep down into the Protocol Buffers core API.

Thanks for the reply. Is there any plan from the Protocol Buffer side to address this issue ? Otherwise, it will be a problem in the future. Also curious, why pytorch does not have this problem ?

[davidel]

I don't know the details but this has been a plague that affected PB for a long time, and if it is still no fixed, I am guessing there are good reasons.
IIRC it was due to the PB API having int all over, as buffer size types, but I am not entirely sure.

[fingoldo]

I am guessing there are good reasons

I am guessing someone is lazy and is doing his job not very responsively (.