Sparsity Made Easy – Introducing the Cerebras PyTorch Sparsity Library

# Construct model and optimizer
model = torch.nn.Module(...)
optimizer = cstorch.optim.SGD(...)

# Model forward and backward
loss = model(...)
loss.backward()

# Update weights.
optimizer.step()

# Construct model and optimizer as usual
model = torch.nn.Module(...)
optimizer = cstorch.optim.SGD(...)

# Construct a sparsity optimizer, and use the returned wrapper as a drop-in replacement for the original optimizer
optimizer = cstorch.sparse.configure_sparsity_wrapper(
   model,
   optimizer,
   sparsity_type=...,
   sparsity=...,
   init_method=...,
)

# Model forward and backward as usual. Sparsity is automatically applied.
loss = model(...)
loss.backward()

# Update weights. If using dynamic sparsity, it is also updated according to its schedule.
optimizer.step()

Base Optimizer: Most state-of-the-art sparsity and pruning algorithms share similar routines for applying and updating masks for models. We consolidate these under a BaseSparsityOptimizer implementation, which handles other standard functions such as defining the checkpoint states for sparsity, initializing masks with custom distributions, and handling the sparse views of the model parameters and optimizer states. This allows users to define new optimizers relatively easily without worrying about the control flow of sparse training and checkpointing.

Update Schedules: Sparse training algorithms often change the mask patterns and sparsity level at different frequencies and rely on some scheduling functions to enable them. We provide a BaseHyperParameter class, which can be used to define custom schedules easily and pass them to existing algorithms. We also implement standard schedules such as cosine, polynomial, and periodic.

Tensor utilities: We provide a few base utilities for handling mask updates for sparsity algorithms, enabled via an efficient TopK implementation for the CS2. To enable fine-grained tensor handling, we also develop two utilities for developers:

1. 1. ScoreShaper: enables parameter reshaping, allowing for grouped structures during training.
   2. ScoreTieBreaker: enables breaking ties between individual elements within a tensor when writing custom sparsification logic to ensure determinism.

We rely on the above-defined abstractions of our API to implement the following sparse training algorithms as a representative set of baselines:

1. 1. Static Sparse Training
   2. Gradual Magnitude Pruning (GMP) [4]
   3. Sparse Evolutionary Training (SET) [1]
   4. Rigging the Lottery (RigL) [2]

Our developer documentation contains more details for functional arguments and support for update schedules, initializations, etc. We also provide a detailed guide on how to set up a sparse training workflow from scratch for training on the Cerebras CS2.

Why support dynamic sparsity?

The Lottery Ticket Hypothesis [5] demonstrated that we can find a sparse network with iterative pruning and successfully train it from scratch to achieve comparable accuracy by starting from the original initial conditions. In practice, this work relies on finding the “winning ticket,” which is compute-intensive and often challenging to discover. Previous works, such as SNIP [6], GRASP [7], etc., have tried finding this winning ticket at initialization to reduce compute costs but lose accuracy compared to training a dense model. As an orthogonal approach, some works, such as SET [1], RigL [2], etc., have focused on employing dynamic updates to efficiently identify optimal sparse networks within a single training cycle, bypassing the need for finding the winning ticket. Figure 3 illustrates the general workflow of dynamic updates during training. The recent state-of-the-art research on sparsity for training neural networks relies on dynamic sparse methods by default. Also, in our recent work, Sparse-IFT [8], we benchmark the advantages of dynamic sparse training over static sparse training and show consistent wins at all sparsity levels.

Push Button Software for Sparse Training

The Cerebras Software Platform makes it extremely simple to pre-train models using unstructured sparsity. Any existing PyTorch model in the Cerebras Model Zoo can be made sparse with just a few lines of change to the configuration file, as shown in Figure 4.

Benchmarks

We benchmark our sparse training algorithms for training LLMs to demonstrate the new API and the effectiveness of dynamic sparsity beyond static sparsity. We use the same architecture as the Llama2 [9] family of models but do not adopt their findings on Generalized Query Attention (GQA) and long context lengths. We train a 1.3 billion (B) parameter model on 112 billion tokens of SlimPajama [10] data. Table 1 shows the architectural details of the model. We use the AdamW [3] optimizer with betas of (0.9, 0.95) and epsilon of 10^-8. The global norm is clipped at 1.0, and a weight decay of 0.1 is used. There is a learning rate warm-up over the first 2000 steps to a peak value of 2 ∗ 10^-4, followed by a cosine decay to 4.5 ∗ 10^-6. We train on packed sequences of 2048 context length for computational efficiency.

To train the sparse models, we uniformly prune 80% of all the linear layer weights in the decoder blocks (5x compression). The normalization, embeddings, and output linear layers are not pruned to promote training stability for the sparse models. Following the findings in the In-Time Overparameterization [11] paper, we reduce the batch size by half compared to the dense model and train for 2x longer. We also increase the drop fraction to 0.5, follow a cosine decay pruning schedule, and decrease the frequency of updates to allow algorithms such as SET and RigL to find better masks. The hyper-parameters for all models are shown in Table 2.

While a single CS-2 system can seamlessly pre-train GPT models up to 175 billion parameters, we leverage a Wafer-Scale Cluster equipped with 4 x CS-2 systems to scale pre-training to speed up our experiments. The remarkable ease of scaling is shown in Figure 5. A more detailed discussion of the CS-2’s scaling properties can be found in this blog post.

class Constant(BaseHyperParameter):

    """
    Constant at every step.
    """
  
  TYPE = "constant"

    def __init__(self, value):
        self.value = torch.tensor(value)

    def __call__(self, step: torch.Tensor, is_update_step: torch.Tensor):
        return self.value

class Cubic(BaseHyperParameter):

    """
    Cubic sparsity function.

    :math:`s_t = s_f + (s_i - s_f) * (1 - (t - t0) / (n * t_delta))**3`
    """

    TYPE = "cubic"

    def __init__(
        self,
        init_sparsity,
        end_sparsity,
        sparsity_start_step,
        sparsity_end_step,
        prune_every_k_steps,
    ):
        self.s_init = init_sparsity
        self.s_end = end_sparsity
        self.update_iter = prune_every_k_steps
        self.init_iter = int(sparsity_start_step / prune_every_k_steps)
        self.final_iter = int(sparsity_end_step / prune_every_k_steps)
        self.total_iters = self.final_prune_iter - self.initial_prune_iter

       def __call__(self, step: torch.Tensor, is_update_step: torch.Tensor):

        curr_iter = (step / self.update_iter).int()
        prune_decay = (1 - ((curr_iter - self.init_iter) / self.total_iters)) ** 3
        current_prune_rate = self.s_end + (self.s_init - self.s_end) * prune_decay
        return torch.clamp(current_prune_rate, min=self.s_init, max=self.s_end)

    def get_min_max_end(
       self, begin: int, end: int
    ) -> Tuple[float, float, float]:
        return (self.init_sparsity, self.end_sparsity, self.end_sparsity)

We train a Llama2 1.3B model (from our benchmarks) following the same training configurations for 2.6B tokens using our GraNet implementation and show the training curves below in Figure 8.

Conclusion

In this blog, we introduce our PyTorch-based library for training models with weight sparsity and show results for training some large models with it. We also show how easy integrating new algorithms and enabling sparsity for training models with the Cerebras Model Zoo is. Parallel to our work, libraries like JaxPruner [16] and STen [17] have also been released to enable sparsity research.

The Cerebras CS-2’s specialized architecture enables unprecedented efficiency and performance for sparse neural network models. Our co-designed ML/software solution allows users to access this performance through a research-friendly API. This library is already pivotal in supporting our in-house research on sparsity, demonstrated through the works in Sparse Pre-training and Dense Fine-tuning [14, 15] and Sparse Iso-FLOP Transformations [8].

We are actively exploring new methods and directions to optimize performance and the quality of sparse models. Contact us to learn more about this study or how the Cerebras CS-2 and our software platform can empower your sparsity research.

References

1. 1. 1. Mocanu, Decebal Constantin, et al. “Scalable training of artificial neural networks with adaptive sparse connectivity inspired by network science.” Nature communications1 (2018): 2383.
      2. Evci, Utku, et al. “Rigging the lottery: Making all tickets winners.” International Conference on Machine Learning. PMLR, 2020.
      3. Loshchilov, Ilya, and Frank Hutter. “Decoupled weight decay regularization.” arXiv preprint arXiv:1711.05101 (2017).
      4. Zhu, Michael, and Suyog Gupta. “To prune, or not to prune: exploring the efficacy of pruning for model compression.” arXiv preprint arXiv:1710.01878 (2017).
      5. Frankle, Jonathan, and Michael Carbin. “The lottery ticket hypothesis: Finding sparse, trainable neural networks.” arXiv preprint arXiv:1803.03635 (2018).
      6. Lee, Namhoon, Thalaiyasingam Ajanthan, and Philip HS Torr. “Snip: Single-shot network pruning based on connection sensitivity.” arXiv preprint arXiv:1810.02340 (2018).
      7. Wang, Chaoqi, Guodong Zhang, and Roger Grosse. “Picking winning tickets before training by preserving gradient flow.” arXiv preprint arXiv:2002.07376 (2020).
      8. Saxena, Shreyas, et al. “Sparse Iso-FLOP Transformations for Maximizing Training Efficiency.” arXiv e-prints (2023): arXiv-2303.
      9. Touvron, Hugo, et al. “Llama 2: Open foundation and fine-tuned chat models, 2023.” URL https://arxiv. org/abs/2307.09288 (2023).
      10. Soboleva, Daria, et al. “SlimPajama: A 627B token cleaned and deduplicated version of RedPajama.” (2023).
      11. Liu, Shiwei, et al. “Do we actually need dense over-parameterization? in-time over-parameterization in sparse training.” International Conference on Machine Learning. PMLR, 2021.
      12. Liu, Shiwei, et al. “Sparse training via boosting pruning plasticity with neuroregeneration.” Advances in Neural Information Processing Systems 34 (2021): 9908-9922.
      13. Frantar, Elias, et al. “Scaling laws for sparsely-connected foundation models.” arXiv preprint arXiv:2309.08520 (2023).
      14. Thangarasa, Vithursan, et al. “SPDF: Sparse Pre-training and Dense Fine-tuning for Large Language Models.” arXiv preprint arXiv:2303.10464 (2023).
      15. Gupta, Abhay, et al. “Accelerating Large Language Model Training with Variable Sparse Pre-training and Dense Fine-tuning.” Cerebras Blog (2023)
      16. Lee, Joo Hyung, et al. “JaxPruner: A concise library for sparsity research.” Conference on Parsimony and Learning. PMLR, 2024.
      17. Ivanov, Andrei, et al. “STen: Productive and Efficient Sparsity in PyTorch.” arXiv preprint arXiv:2304.07613 (2023).

Contributions

Abhay Gupta led the design of the sparsity API along with Mark Browning, tested the algorithms during their internal bring-up, ran all benchmarks, and contributed to the writing of this blog. Mark Browning is the primary developer of the sparsity library, enabling framework support for the CS-2. Claire Zhang helped implement the GraNet algorithm, ran the associated experiments, and contributed to writing the blog. Sean Lie is the architect of sparsity on the Cerebras CS-2 and has guided the bring-up of the hardware, software, and training infrastructure used for the training runs in this blog. We also acknowledge the software, machine learning, and performance teams, who have played an instrumental role in developing sparsity support on the CS-2 hardware.