Overview

Welcome to Sparkformers, where we offer distributed training of Transformers models on Spark!

Motivation / Purpose

Derived from Elephas, however with HuggingFace removing support for Tensorflow, I decided to spin some of the logic off into its own separate project, and also rework the paradigm to support the Torch backend! The purpose of this project is to serve as an experimental backend for distributed training that may be more developergonomic compared to other solutions such as Ray. Additionally, Sparkformers offers the capability for distributed prediction, model calling, and generation (for causal/autoregressive models).

The project is currently in a beta/experimental state. While not yet production ready, I invite you to experiment, provide feedback, and/or even contribute!

Approach

Training: The current architecture utilizes federated averaging (FedAvg), meaning that each executor is trained on a subset of data, and the model weights are averaged across all executors after each epoch. The original model is then updated with the averaged weights, and then the process is repeated for the next epoch.

Inference: The input data is distributed across the executors, and each executor performs the inference on its subset of data. The results are then collected and returned to the driver.

Generation: Same as above, but with the generate method of the model.

Installation

To install, you can simply run: bash pip install sparkformers `

(or uv add, poetry add, etc. with whichever project dependency management tool you may use).

Examples

Note that all examples are also available in the examples directory.

Autoregressive (Causal) Language Model Training and Inference

```python from datasets import loaddataset from sklearn.modelselection import traintestsplit from sparkformers.sparkformer import Sparkformer from transformers import ( AutoTokenizer, AutoModelForCausalLM, ) import torch

batch_size = 16 epochs = 100

dataset = loaddataset("gfigueroa/wikitextprocessed") x = dataset["train"]["text"]

xtrain, xtest = traintestsplit(x, test_size=0.1)

model_name = "hf-internal-testing/tiny-random-gptj"

model = AutoModelForCausalLM.frompretrained(modelname) tokenizer = AutoTokenizer.frompretrained(modelname) tokenizer.padtoken = tokenizer.eostoken tokenizerkwargs = { "maxlength": 50, "padding": True, "truncation": True, "padding_side": "left", }

sparkformermodel = Sparkformer( model=model, tokenizer=tokenizer, loader=AutoModelForCausalLM, optimizerfn=lambda params: torch.optim.AdamW(params, lr=1e-3), tokenizerkwargs=tokenizerkwargs, num_workers=2, )

perform distributed training

sparkformermodel.train(xtrain, epochs=epochs, batchsize=batchsize)

perform distributed generation

generations = sparkformermodel.generate( xtest, maxnewtokens=10, numreturnsequences=1 )

decode the generated texts

generatedtexts = [ tokenizer.decode(output, skipspecial_tokens=True) for output in generations ]

for i, text in enumerate(generatedtexts): print(f"Original text {i}: {xtest[i]}") print(f"Generated text {i}: {text}") ```

Sequence Classification

```python from datasets import loaddataset from sklearn.modelselection import traintestsplit from torch import softmax

from sparkformers.sparkformer import Sparkformer from transformers import ( AutoTokenizer, AutoModelForSequenceClassification, ) import numpy as np import torch

batch_size = 16 epochs = 20

dataset = loaddataset("agnews") x = dataset["train"]["text"][:2000] y = dataset["train"]["label"][:2000]

xtrain, xtest, ytrain, ytest = traintestsplit(x, y, test_size=0.1)

model_name = "prajjwal1/bert-tiny"

model = AutoModelForSequenceClassification.frompretrained( modelname, numlabels=len(np.unique(y)), problemtype="singlelabelclassification", )

tokenizer = AutoTokenizer.frompretrained(modelname) tokenizerkwargs = {"padding": True, "truncation": True, "maxlength": 512}

sparkformermodel = Sparkformer( model=model, tokenizer=tokenizer, loader=AutoModelForSequenceClassification, optimizerfn=lambda params: torch.optim.AdamW(params, lr=2e-4), tokenizerkwargs=tokenizerkwargs, num_workers=2, )

perform distributed training

sparkformermodel.train(xtrain, ytrain, epochs=epochs, batchsize=batch_size)

perform distributed inference

predictions = sparkformermodel.predict(xtest) for i, pred in enumerate(predictions[:10]): probs = softmax(torch.tensor(pred), dim=-1) print(f"Example {i}: probs={probs.numpy()}, predicted={probs.argmax().item()}")

review the predicted labels

print([int(np.argmax(pred)) for pred in predictions]) ```

Token Classification (NER)

```python from sklearn.modelselection import traintestsplit from sparkformers.sparkformer import Sparkformer from transformers import ( AutoTokenizer, AutoModelForTokenClassification, ) from datasets import loaddataset import numpy as np import torch

batchsize = 5 epochs = 1 modelname = "hf-internal-testing/tiny-bert-for-token-classification"

model = AutoModelForTokenClassification.frompretrained(modelname) tokenizer = AutoTokenizer.frompretrained(modelname)

def tokenizeandalignlabels(examples): tokenizedinputs = tokenizer( examples["tokens"], truncation=True, issplitintowords=True ) labels = [] for i, label in enumerate(examples["nertags"]): wordids = tokenizedinputs.wordids(batchindex=i) previouswordidx = None labelids = [] for wordidx in wordids: if wordidx is None: labelids.append(-100) elif wordidx != previouswordidx: labelids.append(label[wordidx]) else: labelids.append(-100) previouswordidx = wordidx labels.append(labelids) tokenizedinputs["labels"] = labels return tokenized_inputs

dataset = loaddataset("conll2003", split="train[:5%]", trustremotecode=True) dataset = dataset.map(tokenizeandalignlabels, batched=True)

x = dataset["tokens"] y = dataset["labels"]

xtrain, xtest, ytrain, ytest = traintestsplit(x, y, test_size=0.1)

tokenizerkwargs = { "padding": True, "truncation": True, "issplitintowords": True, }

sparkformermodel = Sparkformer( model=model, tokenizer=tokenizer, loader=AutoModelForTokenClassification, optimizerfn=lambda params: torch.optim.AdamW(params, lr=5e-5), tokenizerkwargs=tokenizerkwargs, num_workers=2, )

sparkformermodel.train(xtrain, ytrain, epochs=epochs, batchsize=batch_size)

inputs = tokenizer(xtest, **tokenizerkwargs) distributedpreds = sparkformermodel(**inputs) print([int(np.argmax(x)) for x in np.squeeze(distributed_preds)])

```

TODO

• [ ] Add support for distributed training of other model types (e.g., image classification, object detection, etc.)
• [ ] Support training paradigms using Trainer, TrainingArguments, and DataCollater
• [ ] Expose more configuration options
• [ ] Consider simplifying the API further (e.g; builder pattern, providing the model string and push loader logic inside the Sparkformer class, etc.) > 💡 Interested in contributing? Check out the Local Development & Contributions Guide.