SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zachary Charles*
zachcharles@google.com
Nicole Mitchell*
nicolemitchell@google.com
Krishna Pillutla*
kpillutla@google.com
Michael Reneer
michaelrenee@google.com
Zachary Garrett
zachgarrett@google.com
We introduce Dataset Grouper, a library to create large-scale group-structured (e.g., federated) datasets, enabling federated learning simulation at the scale of foundation models. This library facilitates the creation of group-structured versions of existing datasets based on user-specified partitions, and directly leads to a variety of useful heterogeneous datasets that can be plugged into existing software frameworks. Dataset Grouper offers three key advantages. First, it scales to settings where even a single group's dataset is too large to fit in memory. Second, it provides flexibility, both in choosing the base (non-partitioned) dataset and in defining partitions. Finally, it is framework-agnostic. We empirically demonstrate that Dataset Grouper enables large-scale federated language modeling simulations on datasets that are orders of magnitude larger than in previous work, allowing for federated training of language models with hundreds of millions, and even billions, of parameters. Our experimental results show that algorithms like FedAvg operate more as meta-learning methods than as empirical risk minimization methods at this scale, suggesting their utility in downstream personalization and task-specific adaptation. Dataset Grouper is available at https://github.com/google-research/dataset_grouper.
In most machine learning and artificial intelligence settings, algorithms operate on "flat" collections of examples, that is, examples with no differentiation in source or group structure. However, data in the real world often consists of an explicit or implicit underlying group structure, where the examples are partitioned across some number of groups with markedly different statistical characteristics. Increasingly, research has found that this structure is important in a variety of settings, both for encoding constraints (such as data restrictions) and in developing algorithms for learning tasks.
Federated learning (FL) is one such setting. FL methods are designed to operate on data partitioned explicitly across "clients". In cross-device FL [1, Table 1], clients are often edge devices which exhibit heterogeneity in both quantity and distribution of data. Cross-silo FL exhibits a similar structure, often with a coarser notion of clients (such as institutions or companies). Group structures also arise in meta-learning, in which data is grouped according to a notion of "task" [2], and in personalization, in which a user's specific data is used to tune an algorithm's outputs. The group structure can also be highly relevant in the context of differential privacy [3, 4]. An intuitive "unit
Source	Dataset	Group by	Words	Groups	Words per group	Examples	Words per example
10thperc.	Median	90thperc.	10thperc.	Median	90thperc.
Ours	FedC4	Domain	132B	15.6M	82	815	0.5B	0.36B	49	191	783
FedWiki	Article	3B	6.5M	39	198	70K	6.5M	39	198	70K
FedBookCO	Book	1.2B	18K	24K	52K	4M	18K	24K	52K	4M
FedCCnews	Domain	0.3B	8.8K	303	5K	8.4M	0.7M	78	316	842
Existing	Amazon Reviews	Account	4.3B	1.5M	278	1.1K	5K	68M	3	28	155
Stack Overflow	Account	2B	0.3M	1.2K	2.7K	11K	0.1B	3	13	29
Reddit	Account	1.2B	1.7M	58	257	1720	33M	7	21	81
Blog Corpus	Account	0.1B	17K	551	2K	13K	0.5M	6	105	460
Shakespeare	Role/play	0.4M	715	14	175	1.6K	16K	4	12	63
Gigaword	Synthetic	0.3M	100	3.0K	3.1K	3.2K	10K	21	31	41
Table 1: Summary of the per-group (i.e., per-client) and per-example (i.e., per-sequence) statistics of the new language modeling datasets we introduce using Dataset Grouper, compared to those of previous federated benchmark datasets supplied by TFF [11], LEAF [12], FedNLP [13, 14], and FedScale [15].
of privacy" is the total collection of examples associated with a given user [5]. To ensure user-level differential privacy, we must generally train the model in a user-aware manner.
The increasing prominence of foundation models and large language models (LLMs) and their wise applicability to downstream tasks enhance the need for group-structured data. Though foundation models are generally trained on massive flat datasets, they are often evaluated by considering the performance on various benchmarks, yielding a natural group structure. Moreover, for downstream, user-facing applications, one may want to train on user-partitioned data that is representative of the actual task at hand. Alternatively (or in conjunction) one can personalize a foundation model for a given user [6-8]. In all these settings, one may wish to maintain formal user-level privacy guarantees, especially given the privacy and memorization concerns surrounding foundation models [9, 10].
All of the aforementioned scenarios require datasets with explicit group structure. Since foundation models generally require large amounts of data, these research areas may specifically benefit from large-scale group-structured datasets. Unfortunately, to the best of our knowledge, there are relatively few existing datasets that meet such criteria. While a variety of federated datasets are available to researchers [11-15], many of these are small-scale in terms of the number of groups, the quantity of data, or quantity of data per group. Moreover, they may only be available in formats that do not scale well, either due to memory requirements or insufficient efficiency.
Contributions. In this work, we address the growing need for a wider variety of group-structured datasets, especially at larger scales. Concretely, we make the following contributions.
- A library for creating group-structured datasets: We introduce Dataset Grouper, a library that can flexibly create group-structured (and federated) versions of existing datasets via user-defined partition functions. We engineer it for efficiency, both in partitioning datasets and in iterating over data. The library is designed with large-scale datasets in mind and can support datasets with millions or even billions of groups. Dataset Grouper can be used to create group-structured versions of all datasets available in Tensorflow Datasets [16] and HuggingFace Datasets [17].
- Large-scale federated text datasets: While Dataset Grouper can be used for a wide array of modalities and tasks, we illustrate its use by creating group-structured versions of four large language modeling datasets (see Table 1 and Figure 1), designed specifically for FL research. They are orders of magnitude larger than previous datasets in terms of one or more of the following: the number of groups, the amount of data, and the length of sequences. They are suitable for both pre-training and fine-tuning tasks, and exhibit long tails, as is common in large-scale text corpora.
- Experiments: We train  $O(100\mathrm{M})$  and  $O(1\mathrm{B})$  parameter decoder-only transformer models from scratch on a group-structured version of the C4 dataset [18], using FL training algorithms. This is, to the best of our knowledge, the first demonstration of federated training of a model of this magnitude on a federated dataset of this scale. We compare FedAvg and FedSGD in this setting in terms of their pre- and post-personalization metrics. Our results highlight that at this scale, FedAvg behaves more like a meta-learning algorithm than an empirical risk minimization algorithm.
On the term "federated". Our work is primarily motivated by the research needs of the FL community. Throughout, we will often approach questions, design decisions, and experiments from this
Figure 1: Per-group statistics of the new group-structured (i.e. federated) language modeling datasets.
perspective, and will occasionally use the terms "federated" and "group-structured" interchangeably. This is a broadening of the definition of the federated setting often given in the literature, especially the definition proposed by Kairouz et al. [19]. There and in previous literature, FL is characterized by the group-level structure of the data coupled with the location of each group (e.g., a client's dataset is assumed to reside only on its local device). In this work, we de-emphasize the location of the data, and will primarily focus on the group-structure of the data regardless of where that data lives.
Many widely used group-structured datasets arise from the FL community. Early work on FL benchmark datasets combined benchmark datasets with simulation environments. For example, LEAF [12], TensorFlow Federated (TFF) [11], Flower [20], FedML [14] (and its text-specific dataset provided in FedNLP [13]), FedScale [15], FLBench [21], OARF [22], and FedJAX [23] all supply partitioned datasets commonly used in FL benchmarks. These frameworks have helped drive much of FL research into areas such as optimization algorithms [24-31], privacy and security [32-35], robustness to adversaries [36-40] and distribution shifts [41-45], and personalization [46-51].
Later work on FL introduced specialized datasets and benchmarks. FLAIR [52] is a large-scale multi-label image classification dataset with 0.4M images from 50K Flickr accounts. FLamby [53] is a cross-silo FL benchmark for healthcare applications where the datasets contain 2-6 groups with 400 to 23K total examples. The personalized federated learning algorithm benchmark Motley [54]'s largest dataset is Stack Overflow while pFL-bench [55] offers no language modeling datasets.
The scale of the existing language modeling datasets that FL frameworks provide is summarized in Table 1. By comparison, the datasets we provide are significantly larger, which allows training models that are an order of magnitude larger than previous work [56]. Additionally, our datasets are generally framework-agnostic and can be combined with many of these simulation frameworks. Further, existing federated datasets that are derived from datasets in TensorFlow Datasets [16], such as Amazon Reviews and Blog Corpus, can be generated via Dataset Grouper.
Group-structured data have also been studied in the context of distribution shifts, e.g., the WILDS benchmark [57, 58]. The datasets provided in WILDS are smaller than what we consider in Table 1 — the code language modeling dataset Py150 dataset has  $150\mathrm{K}$  examples partitioned over  $8\mathrm{K}$  groups. Moreover, WILDS does not support optimized per-group data-loaders as of this writing [59], which are necessary for benchmarking federated algorithms.
LLMs are typically pre-trained on large web-scraped text corpora without any group structure [e.g. 18, 60, 61]. Besides the tremendous amount of data on which they are trained [62], the success of LLMs is also driven by the capacity of these models to handle much longer sequences than previous RNN-based models [63-65]. This requires datasets with long enough contiguous sequences that contain hundreds to thousands of words. Almost all of the existing group-structured language modeling datasets have extremely short sequences (Table 1). For instance, the Stack Overflow dataset has a median and  $90^{\text{th}}$  percentile sequence lengths of 13 and 29 words respectively. In comparison, the datasets we introduce have significantly longer sequences, e.g., FedBookCO has on the order of  $10^{3}$  to  $10^{6}$  words per sequence.
Some recent advances at the intersection of FL and foundation models include collaborative prompt tuning using FL [66], federating chain-of-though reasoning [67] through diverse crowd-workers [68], and instruction-tuning LLMs using FL [69]. Dataset Grouper can also be used to generate federated datasets compatible with these methods as well.
(a) In-Memory.
(b) Hierarchical.
(c) Streaming.
Figure 2: High-level representations of group-structured dataset formats.
3 Core Design
We now discuss the core design of Dataset Grouper and the various trade-offs it entails. One unifying theme throughout is a focus on enabling the types of training workflows used for foundation models, at the expense of some amount of flexibility in what kinds of simulations can be performed.
3.1 Group-Structured Datasets at Scale
Our primary goal is to enable research using large-scale group-structured datasets. In order to do so, we need a group-structured dataset format that balances the following characteristics:
- Scalability: Can the dataset format scale to large numbers of examples, groups, and bytes?
- Group access time: How long does it take to access the examples held by a single group?
- Group access patterns: What kinds of sampling patterns across groups are permitted? Can we access group datasets arbitrarily, and in any order?
There is a trade-off between these characteristics. Dataset formats used in the FL community often optimize for either scalability or group sampling time, while enabling maximum flexibility in access patterns. Our core insight is that by limiting the access patterns possible, we can use a dataset format that is scalable and efficient simultaneously. We discuss three archetypes of group-structured dataset formats (in-memory, hierarchical, and streaming) and their resulting trade-offs briefly in Table 2, and in more detail below. Figure 2 gives a graphical representation of the formats.
In-memory formats. In-memory group-structured datasets are essentially key-value mappings held entirely in memory. Adopted by e.g. LEAF [12] and FedNLP [13], this is suitable for small datasets such as EMNIST or CIFAR-100. Looking up a group's dataset is fast (e.g., via a hash map), and groups can be accessed in an arbitrary order. Of course, this approach is limited by the cumulative size of the dataset and is therefore not scalable in full generality. As we see in Table 3, this format does not even scale to FedBookCO on a single CPU; FedC4 and FedWiki are even larger.
Hierarchical formats. Hierarchical dataset formats store examples across files in such a way that (a) the dataset need not be loaded entirely into memory, and (b) individual groups can be constructed in arbitrary orders. For example, TensorFlow Federated [11] uses SQL databases to both store and access client datasets for FL simulations, facilitating the loading of the group index in-memory, then construction of a group's dataset at a later time. For larger datasets, constructing an arbitrary group's dataset can be slow, as it is often bottlenecked by indexing and searching over a large number of (possibly distributed) files. Table 3 shows that the hierarchical format can be significantly slower than other formats when accessing groups in very large datasets.
In-Memory	Hierarchical	Streaming
Scalability	Limited	High	High
Group Access Time	Very Fast	Slow	Fast
Group Access Patterns	Arbitrary	Arbitrary	Shuffle + Streaming
Table 2: Characteristics of group-structured dataset formats.
Dataset Format	In-Memory	Hierarchical	Streaming
CIFAR-100	0.0783 ± 0.0007	25.11 ± 0.81	9.88 ± 0.075
FedCCnews	0.549 ± 0.014	>7200	248 ± 17.5
FedBookCO	Out of memory	>7200	192 ± 9.07
Table 3: The time (in seconds) to iterate over federated datasets. This is the time required to iterate over all examples in all group datasets, in serial, on a single CPU. We present the average and standard deviation over 5 trials, omitting trials that take more than 2 hours ( $>7200$  seconds), or that ran out of memory. We compare a federated CIFAR-100 dataset (partitioned across 100 groups, each with 100 examples), FedCCnews (in which examples are split across users at a domain level), and FedBookCO (in which examples are split across users at a title level). See Section 4 for more details on the latter two datasets.
Streaming formats. Instead of allowing arbitrary group access, Dataset Grouper provides ways to iterate over all the groups in a stream-like fashion. The datasets for each group are backed by some number of files, $^{2}$  which are interleaved to create a "group stream". Concretely, this restricts the possible group access patterns, only allowing stream-level operations such as buffered shuffling, repeating, and batching. This essentially lifts the stream-of-examples format used large-scale centralized training pipelines to streams of groups for federated training — both formats allow dataset iterators with limited shuffling (e.g., with a fixed-size buffer), but not arbitrary access to the individual elements. This restriction allows us to use parallel reads, prefetching, and interleaving to speed up dataset iteration and generally enables the total iteration time of the dataset to scale linearly (as opposed to super-linearly) with the number of groups in the dataset.
Each group's dataset is further represented as a stream of examples so that no group's data need to be fully loaded into memory. This is crucial in scaling to large datasets like FedC4, something that is memory-prohibitive for in-memory formats, and speed-prohibitive for hierarchical formats. To illustrate this further, we detail the time required to iterate fully over various group-structured datasets (accessing the groups' datasets sequentially, in a random order) in different formats in Table 3. For details on the amount of memory used by each format, see Appendix E.
3.2 Flexible and Efficient Dataset Partitioning
An underlying theme of both foundation model research and FL research is the need for a wide variety of datasets. It is often useful to have different datasets for different downstream tasks and modalities for foundation models, while the wide variety of FL settings (e.g. cross-device vs. cross-silo) and types of group heterogeneity (feature heterogeneity, label heterogeneity, heteroskedasticity, etc.) require dedicated datasets. It is often useful in FL to be able to explicitly partition the same dataset in multiple ways, in order to understand the impact of heterogeneity [71]. Therefore, our second key design goal is to allow flexibility both in the base (non-partitioned) dataset and in how it is partitioned.
To achieve this, we make two important, albeit related, design decisions. First, Dataset Grouper does not directly host datasets, but instead allows users to create partitioned versions of existing datasets in TensorFlow Datasets [16] and HuggingFace Datasets [17]. Second, Dataset Grouper operates by applying data-parallel processing pipelines<sup>3</sup> to partition these "base" datasets. Notably, Dataset Grouper allows user-specified partition methods, but they must operate in an embarrassingly parallel manner. This decision is a formal trade-off made for scalability reasons. Sequential partitioning (e.g., deciding which group has an example  $x$  based on which group has an example  $y$ ) can fail to scale to datasets with billions of examples. Thus, Dataset Grouper supports (at scale) embarrassingly parallelizable partitions of datasets available in TensorFlow Datasets or HuggingFace Datasets.
3.3 Compatibility with Existing Frameworks
Foundation model research and FL research also share a common feature in that there is a wide array of available simulation frameworks. Another goal of our work is to support as wide an array of such frameworks as possible. To that end, Dataset Grouper provides access to datasets as nested iterators
2We use the TFRecord format [70] for all datasets.
3We use Apache Beam pipelines, which are also used by TensorFlow Datasets and HuggingFace Datasets.
Figure 3: Fitting a log-normal distribution to the per-group sizes of the new text datasets we introduce: we show a Q-Q plot of the log quantiles of the per-group data sizes vs. those of a Gaussian distribution.
of tensors. Specifically, group-structured datasets are represented as an iterator of group datasets, each of which is an iterator of tensors. These tensors can be represented in both TensorFlow [70] and NumPy [72] formats, ensuring that, in principle, Dataset Grouper can be used in any simulation framework built on top of NumPy, TensorFlow, PyTorch [73], or JAX [74].<sup>4</sup>
4 Examples and Applications
We now focus on four new group-structured text datasets we create via Dataset Grouper: FedC4, FedWiki, FedCCnews, FedBookCO. We focus on language modeling datasets due to their prominence in training foundation models and their large-scale nature. Compared to prior benchmark datasets, FedC4 is an order of magnitude larger, while FedBookCO contains significantly longer sequences. The new datasets, particularly FedC4 and FedCCnews, are also more heavy-tailed than existing ones. See Appendix B for more details.
While representative of the statistical structure suited to training larger models, we wish to emphasize that these datasets are only a small sample of what is possible with Dataset Grouper. The library can also be used to create group-structured multi-lingual text datasets, datasets in other modalities (audio, image, etc.), and to study the effect of different partitions on the same base dataset.
FedC4. We create a federated version of the Colossal Clean Crawled Corpus (C4) dataset [18], which is a cleaned version of Common Crawl's web crawl corpus.<sup>5</sup> We focus on partitioning by web domain, e.g., all articles crawled from https://www.nytimes.com/ correspond to one group. We note that a finer partitioning at the level of articles is also possible. We see from Figure 1 and Table 1 that the amount of data per client is extremely heavy-tailed; this is expected from real-world text corpora [75, 76]. Indeed, this distribution is nearly log-normal, meaning that its logarithm is approximately Gaussian. This can be seen from the nearly straight line in the Q-Q plot in Figure 3.
The C4 data is also used as a pre-training corpus for some LLMs such as T5 [18], meaning that FedC4 can potentially be used for federated pre-training, which we explore further in Section 5. Note that C4 is already de-duplicated and artifacts like code, placeholder text (e.g. lorem ipsum), and boilerplate (e.g. menus) are removed along with a heuristic list of "bad words". See [18] for details.
FedWiki. We create a federated version of the Wikipedia dataset, where each client contains the content of one full English Wikipedia article. As a result, the amount of data per client is smaller than that in FedC4, where each client can contain multiple articles. Wikipedia data is often a part of the pre-training corpora of LLMs.
FedBookCO. We create a federated version of the BookCorpusOpen dataset [77, 78], an open-source collection of 18K books from various genres. Each client corresponds to one sequence that is a full book, leading to significantly longer sequences than other datasets.
FedCCnews. We create a federated version of CC-News, which contains English news articles from news sites around the world. Similar to FedC4, each group corresponds to a web domain; a finer-grained article-level partitioning is also possible. Indeed, FedCCnews is a subset of FedC4. It exhibits similar long-tailed behavior and could serve well as its smaller proxy.
<sup>4</sup>Some frameworks are only inter-operable with specific in-memory or hierarchical dataset formats, and would need to be extended to be compatible with Dataset Grouper. Other frameworks are directly compatible. We provide examples of integrating Dataset Grouper with both TensorFlow Federated and JAX. <sup>5</sup>https://commoncrawl.org/
Cohort Size	Data Iteration Time (s)	Training Time (s)	Data Iteration Time (%)
8	0.26 ± 0.48	3.03 ± 2.58	7.78
16	0.66 ± 0.85	5.70 ± 2.61	10.43
32	1.16 ± 1.48	11.30 ± 2.48	9.33
Table 4: Average time spent per round on iterating over data, including preprocessing, versus training. Results are computed for 100 rounds of training of FedAvg, with varying cohort sizes.
To begin to demonstrate the scale of federated learning simulation that these newly partitioned datasets enable, we run experiments on FedC4 with a decoder-only transformer architecture.
We use FedC4 with domain-level partitioning in our experiments. We use a WordPiece tokenizer [79] with a pre-trained BERT vocabulary [80] of size of 30523. We train from scratch a 108M parameter decoder-only transformer model commensurate in size with BERT base and GPT-2 small (i.e., 12 layers, 12 attention heads, and hidden layers of dimension 768) using the causal language modeling loss (i.e., next token prediction with cross-entropy). We report the cross-entropy loss throughout, which equals the logarithm of the perplexity.