def test_fast_forward_sampler_over_iterative_dataset(num_workers):
"""
This test ensures ``FastForwardSampler`` and ``CaptureIterableDataset`` are properly being
used to capture workers states.
"""
batch_size = 3
initial_seed = seed_everything(42)
generator = torch.Generator()
generator.manual_seed(initial_seed)
dataset = RangeIterableDataset(range(20), num_workers, batch_size, True)
dataset = CaptureIterableDataset(dataset)
dataloader = DataLoader(dataset,
batch_size=batch_size,
num_workers=num_workers,
generator=generator)
_add_capture_metadata_collate(dataloader)
iter_dataloader = iter(dataloader)
batches = []
for _ in range(5):
batches.append(next(iter_dataloader))
# restarting on batch_1 and getting 3 extra batches
state_dict = {"iter_sampler": {}}
for batch in batches[:2]:
batch, _state_dict = CaptureIterableDataset.extract_samplers_state_dict_from_batch(
batch)
for k, v in _state_dict[0].items():
state_dict[k].update(v)
assert len(
state_dict["iter_sampler"]) == (num_workers if num_workers > 1 else 1)
initial_seed = seed_everything(42)
generator.manual_seed(initial_seed)
dataset = RangeIterableDataset(range(20),
num_workers,
batch_size,
state_dict=state_dict)
dataset = CaptureIterableDataset(dataset)
dataset.load_state_dict(state_dict)
dataloader = DataLoader(dataset,
batch_size=batch_size,
num_workers=num_workers,
generator=generator)
_add_capture_metadata_collate(dataloader)
iter_dataloader = iter(dataloader)
batches_restart = []
for _ in range(3):
batches_restart.append(next(iter_dataloader))
assert torch.equal(batches_restart[0]["data"], batches[2]["data"])
assert torch.equal(batches_restart[1]["data"], batches[3]["data"])
assert torch.equal(batches_restart[2]["data"], batches[4]["data"])
def next_fn(iterator: Iterator):
batch = next(iterator)
if not _fault_tolerant_enabled():
return batch
# when fault tolerant is enabled, the iterator will return
# `FastForwardSampler` state_dict metadata
# along side with the user data.
# the metadata are extracted and store directly on the iterator
# to simplify the collection on `state_dict` call.
batch, samplers_state_dict = CaptureIterableDataset.extract_samplers_state_dict_from_batch(batch)
# store the `sampler_state_dict` on the iterator
CaptureIterableDataset.store_samplers_state_dict(iterator, samplers_state_dict)
return batch
def _test_fast_forward_sampler_with_distributed_sampler_and_iterative_dataset(rank, worldsize):
if worldsize > 1:
_setup_ddp(rank, worldsize)
def all_gather(tensor, world_size):
tensor_list = [torch.zeros_like(tensor, dtype=torch.int64) for _ in range(world_size)]
torch.distributed.all_gather(tensor_list, tensor)
return tensor_list
initial_seed = seed_everything(42)
generator = torch.Generator()
generator.manual_seed(initial_seed)
num_workers = 2
batch_size = 4
dataset_length = 60
num_classes = 10
labels = np.random.randint(0, num_classes, dataset_length)
dataset = ClassificationDataset(range(dataset_length), labels)
dataset = MetaLearningDataset(
dataset,
batch_size=batch_size,
drop_last=True,
num_workers=num_workers,
global_rank=rank,
world_size=worldsize,
initial_seed=initial_seed,
debugging=True,
shuffle=True,
)
dataset = CaptureIterableDataset(dataset, initial_seed=initial_seed)
dataloader = DataLoader(dataset, num_workers=num_workers, batch_size=1, generator=generator)
Trainer._add_sampler_metadata_collate(dataloader)
epoch_results = []
for _ in range(2):
iter_dataloader = iter(dataloader)
batches = []
while True:
try:
batches.append(next(iter_dataloader))
except StopIteration:
break
epoch_results.append(batches)
dataloader.dataset.dataset.current_task_iteration += 1
assert len(epoch_results) == 2
assert len(epoch_results[0]) == math.ceil((dataset_length / (num_workers * worldsize)) / batch_size) + 2
if worldsize == 1:
assert epoch_results[0][0]["data"]["task_length"] == epoch_results[0][1]["data"]["task_length"]
assert torch.equal(
epoch_results[0][0]["data"]["selected_indexes"], epoch_results[0][1]["data"]["selected_indexes"]
)
assert 0 in epoch_results[0][2][AutoRestartBatchKeys.PL_SAMPLERS]["iter_sampler"] # worker id 0
assert 1 in epoch_results[0][3][AutoRestartBatchKeys.PL_SAMPLERS]["iter_sampler"] # worker id 1
assert not torch.equal(epoch_results[0][2]["data"][0], epoch_results[0][3]["data"][0])
else:
first_task_metadata = all_gather(epoch_results[0][0]["data"]["task_length"], worldsize)
second_task_metadata = all_gather(epoch_results[0][1]["data"]["task_length"], worldsize)
assert torch.equal(first_task_metadata[0], first_task_metadata[1])
assert torch.equal(second_task_metadata[0], second_task_metadata[1])
assert torch.equal(first_task_metadata[0], second_task_metadata[1])
first_batch_list = all_gather(epoch_results[0][2]["data"][0], worldsize)
assert not torch.equal(first_batch_list[0], first_batch_list[1])
second_batch_list = all_gather(epoch_results[0][3]["data"][0], worldsize)
assert not torch.equal(second_batch_list[0], second_batch_list[1])
# restarting on epoch 0 / real batch 2
state_dict = {"iter_sampler": {}}
for batch in epoch_results[0][2:4]:
batch, _state_dict = CaptureIterableDataset.extract_samplers_state_dict_from_batch(batch)
for k, v in _state_dict[0].items():
state_dict[k].update(v)
dataset = ClassificationDataset(range(dataset_length), labels)
dataset = MetaLearningDataset(
dataset,
batch_size=batch_size,
drop_last=True,
num_workers=num_workers,
global_rank=rank,
world_size=worldsize,
initial_seed=initial_seed,
debugging=True,
shuffle=True,
)
dataset = CaptureIterableDataset(dataset, initial_seed=initial_seed)
dataset.load_state_dict(state_dict)
dataloader = DataLoader(dataset, num_workers=num_workers, batch_size=1, generator=generator)
Trainer._add_sampler_metadata_collate(dataloader)
epoch_results_restart = []
for _ in range(2):
iter_dataloader = iter(dataloader)
batches = []
while True:
try:
batches.append(next(iter_dataloader))
except StopIteration:
break
epoch_results_restart.append(batches)
dataloader.dataset.dataset.increment_iteration()
dataloader.dataset.reset_on_epoch()
assert len(epoch_results_restart[0]) + 2 == len(epoch_results[0])
epoch_tensors = [e["data"][0] for e in epoch_results[0][4:]]
epoch_tensors_restart = [e["data"][0] for e in epoch_results_restart[0][2:]]
for t, tr in zip(epoch_tensors, epoch_tensors_restart):
assert torch.equal(t, tr)
epoch_tensors = [e["data"][0] for e in epoch_results[1][2:]]
epoch_tensors_restart = [e["data"][0] for e in epoch_results_restart[1][2:]]
for t, tr in zip(epoch_tensors, epoch_tensors_restart):
assert torch.equal(t, tr)