Include Option to Keep Global Batch Size Constant

[WarrenZhu050413]

Maintain Constant Global Batch Size Upon Failure

With the current implementation of DistributedSampler, the global_batch_size is group_batch_size * num_replica_group.

It may be more preferable if the DistributedSampler is implemented such that the global batch size stays constant. The group_batch_size is global_batch_size//num_replica_group. If it is not completely divisible, then the remaining samples is distributed amongst the available replica groups.

For hardware efficiency reasons one may potentially want to keep the local_batch_size a multiple of a certain integer. If so, the global_batch_size should still be as close to the targetted size as possible, subject to these constraints.

Implementation

The dataloader control logic could be folded into quorum.

The Lighthouse would then coordinate the loading, specifying the local_batch_size along with the start_index for each replica group.

To make this implementation efficient, each replica group do not need to self.wait_quorum() before loading a batch. Instead, since the local batch size could only increase (upon failure), the batch can be loaded assuming num_replica_group remains the same as the last iteration, loading additional samples if needed after receiving quorum.

[WarrenZhu050413]

Here are some implementation details that I worked out today. Would love to hear whether this is suitable to be included in torchFT!

ManagedDistributedSampler

To implement the fault tolerance sampling logic, we need to implement a custom sampler class whose __iter__ method will be a generator that yields a list of dataset indices for the current worker's mini-batch. Since the sampler would need information from the Manager to determine the local batch size and sampling indices, we can follow the pattern in torchFT and implement a ManagedDistributedSampler.

from typing import Iterator, List, Optional, Sized

class ManagedDistributedSampler:
    def __init__(self, manager: Manager, dataset: data.Dataset):
        self.manager = manager
        self.dataset = dataset
        # Implement other logic as needed

    def __iter__(self) -> Iterator[List[int]]:
        """
        Yields per-worker mini-batches of dataset indices for one epoch,
        maintaining a global batch size close to `target_global_batch_size`.

        Per iteration, this method will:
        1.  Fetch `DataLoadingInfo` from the `Manager`.
        2.  Calculate the target local batch size for this worker. This involves:
            a.  Distributing `target_global_batch_size` among `num_replica_groups`.
            b.  Further distributing the replica group's batch size among its `intra_group_world_size` workers.
            c.  Optionally, adjusting the local batch size to be a multiple of `local_batch_size_multiple`,
                while striving to keep the effective global batch size near the target.
        3.  Construct the local mini-batch by:
            a.  Prioritizing indices from `missing_indices` assigned to this worker (based on `global_rank`).
            b.  Filling the remainder with new indices starting from `base_index`, distributed
                across all workers in a round-robin fashion (based on `global_rank` and `global_world_size`).
        4.  Yield the assembled list of indices for the current worker.
        5.  The iteration continues until all dataset samples for the epoch are processed
            and all `missing_indices` (from previous iterations within the epoch) are cleared.
            The Lighthouse will signal when the epoch is complete.
        """
        # Initialize dataset indices (e.g., shuffle if self.shuffle is True)
        # Loop until Lighthouse signals epoch completion (or all data processed)
        #   data_info = self.manager.get_data_loading_info()
        #   ... implement logic from steps 2-4 ...
        #   yield current_worker_indices
        pass # Placeholder for actual implementation

    # Other necessary methods...

get_data_loading_info()

Here, DataLoadingInfo is a dataclass that contains the information needed to load the data correctly for the current worker in the presence of faults. To fetch the DataLoadingInfo from the Manager, we can implement a get_data_loading_info() method to the Manager class.

To have a centralized method getting all the necessary data for the dataloader, we can implement a get_data_loading_info() method to the Manager class. This DataLoadingInfo object, providing each worker with the necessary information to load its portion of the data for the current iteration.

from dataclasses import dataclass
from typing import List, Protocol

@dataclass
class DataLoadingInfo:
    """
    Contains information for a worker to determine its data shard for an iteration.
    """
    # The starting index in the (globally shuffled) dataset for new samples this round.
    base_index: int
    # List of sample indices missed in previous rounds that require reprocessing.
    missing_indices: List[int]
    # Number of currently active replica groups.
    num_replica_groups: int
    # Rank of this worker's replica group among all active replica groups.
    replica_group_rank: int
    # Number of active workers within this worker's replica group.
    intra_group_world_size: int
    # Rank of this worker within its own replica group.
    intra_group_rank: int
    # Total number of active workers across all replica groups.
    global_world_size: int
    # Unique rank of this worker among all active workers globally.
    global_rank: int

class Manager(Protocol): # Using Protocol for interface definition
    """
    Interface for the Manager responsible for coordinating with Lighthouse
    and providing data loading information to workers.
    """
    # ... other Manager methods ...

    def get_data_loading_info(self) -> DataLoadingInfo:
        """
        Retrieves data sharding information for the current iteration.
        This information is sourced from the Lighthouse.
        """
        # Implementation would involve communication with Lighthouse
        # to get the centrally calculated DataLoadingInfo.
        pass

Centralized Data Loading Info Management in Lighthouse

The critical variables base_index and missing_indices will be calculated and maintained centrally within the Lighthouse.

Rationale:

• Consistency: Centralized management guarantees that all replica groups and workers operate on a consistent view of the data distribution, preventing duplicate processing or missed samples.
• Simplified Recovery: This design obviates the need for complex state transfer logic for DataLoadingInfo between Manager instances during Live Checkpoint Recovery. The Lighthouse, as the central coordinator, already maintains the necessary global state.

Mechanism:

1. Initialization: At the start of each epoch, base_index is reset (e.g., to 0), and missing_indices is cleared.
2. Iteration Update:
   • The Lighthouse determines the set of active replica groups for the current iteration (quorum).
   • It calculates the assignment of samples (both new and missing) to workers for the current global batch.
   • Based on this assignment, it updates base_index to point to the next unassigned fresh sample.
3. Failure Handling:
   • When a failure occurs and the quorum_id (or equivalent identifier for the set of active groups) changes, the Lighthouse identifies samples that were assigned to the failed/departed workers in the previous iteration but not yet processed (or whose processing cannot be confirmed). These become missing_indices.
   • These missing_indices are then prioritized in subsequent batch assignments.

By having the Lighthouse manage the global sample assignment strategy, it can dynamically adjust which samples constitute the base_index and which are added to missing_indices, effectively updating the state for the DataLoadingInfo that will be served to Manager instances. The Manager instances then simply fetch this pre-computed information.

[d4l3k]

@WarrenZhu050413 I think this is a totally reasonable approach -- honestly this could be quite useful in a number of places (including for avoiding stragglers). Would be a great feature to have in FT

We do have a fixed batch size mode here: https://github.com/pytorch/torchft/blob/main/torchft/manager.py#L57-L72 but that wastes compute by using hot spares

The one issue with the Sampler approach is that the dataloaders often run in child processes asynchronously so we may need it to be a wrapper around the dataloader to rebatch each step rather than change the sampling paradigm otherwise the batch sizes would lag

There's also some consideration on memory usage which may make this trickier to implement in practice since often you want as large of a batch size that will fit in memory. Gradient accumulation could also be used to maintain a consistent global batch size

[d4l3k]

One other tricky bit here is that with async quorum, we don't actually know the world size until after the forwards pass has started. We could potentially selectively sync if we've detected failures from the previous round so we opportunistically go for the previous batch size when no failures have occurred

[WarrenZhu050413]

Thank you! These are really great feedback with aspects that I didn't take into account.