TPU v3-8 deadlocks when using datasets larger than 2^15 on 8 devices · Lightning-AI/pytorch-lightning#18176

Métricas do repositório

Stars: (26.687 stars)
Métricas de merge de PR: (Mesclagem média 9d 15h) (3 fundiu PRs em 30d)

Description

Bug description

Our tpu v3-8 deadlocks when using multiple 8 TPU cores on large datasets. Specifically, datasets larger than 2^15; one size larger and we get deadlock.

The deadlock occurs somewhere between somewhere between line 222 and line 235 in fit_loop.py. Works fine when using only a single core.

Using lightning 2.0.6 | tpu-vm-pt-2.0

What version are you seeing the problem on?

v2.0

How to reproduce the bug

train_dataloader = DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.num_workers,
        #pin_memory=True,
        #persistent_workers=True,
        drop_last=True,
    )
...

trainer = L.Trainer(
        #callbacks=checkpoint_callbacks,
        #logger=wandb_logger,
        max_epochs=args.epochs,
        accelerator=args.accelerator,
        devices=args.devices,
        precision=args.precision,
        num_sanity_val_steps=0,
        log_every_n_steps=10,
        check_val_every_n_epoch=args.check_val_every_n_epoch,
        #profiler="xla",
        strategy="xla",
    )
    trainer.fit(
        model,
        train_dataloaders=train_dataloader,
        #val_dataloaders=val_dataloader,
    )

Error messages and logs

E0727 11:55:45.723604 1727432 coredump_hook.cc:414] RAW: Remote crash data gathering hook invoked. E0727 11:55:45.723633 1727432 coredump_hook.cc:453] RAW: Skipping coredump since rlimit was 0 at process start. E0727 11:55:45.723659 1727432 client.cc:278] RAW: Coroner client retries enabled (b/136286901), will retry for up to 30 sec. E0727 11:55:45.723670 1727432 coredump_hook.cc:512] RAW: Sending fingerprint to remote end. E0727 11:55:45.723677 1727432 coredump_socket.cc:120] RAW: Stat failed errno=2 on socket /var/google/services/logmanagerd/remote_coredump.socket E0727 11:55:45.723692 1727432 coredump_hook.cc:518] RAW: Cannot send fingerprint to Coroner: [NOT_FOUND] Missing crash reporting socket. Is the listener running? E0727 11:55:45.723700 1727432 coredump_hook.cc:580] RAW: Dumping core locally. E0727 11:55:46.275513 1727432 process_state.cc:784] RAW: Raising signal 11 with default behavior

Environment

No response

More info