# -*- coding: utf-8 -*-
from setuptools import setup

packages = \
['sophia']

package_data = \
{'': ['*']}

install_requires = \
['datasets', 'torch', 'transformers']

setup_kwargs = {
    'name': 'sophia-optimizer',
    'version': '0.2.3',
    'description': 'Sophia Optimizer ULTRA FAST',
    'long_description': '[![Multi-Modality](agorabanner.png)](https://discord.gg/qUtxnK2NMf)\n\n# Sophia Optimizer\n\n[Sophia: A Scalable Stochastic Second-order Optimizer for\nLanguage Model Pre-training](https://arxiv.org/pdf/2305.14342.pdf)\n\nCut Model Training Cost by 50%? with this all-new simple plug in and play Optimizer: Sophia\n\n\n# Usage\n\nDownload with pip ```pip install Sophia-Optimizer``` \n\n\n```python \nimport torch \nfrom torch import nn\nfrom Sophia import SophiaG \n\n\n#or decoupled\n#from Sophia import DecoupledSophia #plug in and play, personalize, improved maintainability\n\n\n#define your model\n\nclass MyModel(nn.Module):\n    def __init__(self):\n        super(MyModel, self).__init__()\n        self.fc1 = nn.Linear(784, 128)\n        self.fc2 = nn.Linear(128, 10)\n\n    def forward(self, x):\n        x = torch.relu(self.fc1(x))\n        x = self.fc2(x)\n        return x\n\n\n#init model loss function and input data\nmodel = MyModel()\nloss_function = nn.CrossEntropy()\ninput_data = ... #input data\n\n#init the optimizer\noptimizer = SophiaG(model.parameters(), lr=2e-4, betas=(0.965, 0.99), rho = 0.01, weight_decay=1e-1)\n\n#decoupled\n#optimizer = DecoupledSophia(model.parameters(), lr=1e-3, betas=(0.9, 0.999), rho=0.04, weight_decay=1e-1,     estimator="Hutchinson")\n\n#training loop\nfor epoch in range(epochs):\n    for batch in data_loader:\n        optimizer.zero_grad()\n        output = model(batch)\n        loss = loss_function(output, target)\n        loss.backward()\n        optimizer.step()\n```\n\n## Training:\nTo run training use git clone method\n\n navigate to experiments folder \n \n ```cd Sophia```\n ```cd experiments```\n\n then run file\n ```python3 training.py````\n\n and if not then do the following:\n\n```python\nfrom Sophia import DecoupledSophia, trainer\n\n\n#train model\ntrainer.train()\n\n\n#eval the model\neval_results = trainer.evaluate()\nprint(f"Perplexity: {torch.exp(torch.tensor(eval_results[\'eval_loss\']))}")\n\n```\n\n\n\nNow with training file ready in experiments folder! 🔥🔥🔥\n\nSophia is an second order clipped stochastic optimization algorithm that uses an inexpensive stochastic estimate of the diagonal of the Hessian as an pre-conditioner and a clipping mechanism to control the worst case update size. It achieves better performance than adam in terms of validation pre-traing loss, total compute, and wall-clock time. By cutting model training cost in half, Sophia can help save millions if not billions of dollars in computational resources.\n\n\n## Benefits\n\nSophia achievs the same validation pre training loss with 50% fewer number of steps than Adam\n\n50% less total compute and 50% less wall-clock time\n\nSeamless integration into existing training pipelines -- plug in and play!\n\nNo special requirments on model architecture or computing infrastructure\n\nSupports both Hutchinson and Gauss-Newton-Bartlett Hessian Estimators\n\n\n\n# Algorithmic pseudocode:\n\n```\n\nInitialize parameters: θ1, learning rate {ηt}, hyperparameters λ, β1, β2, ϵ, and estimator choice Estimator ∈ {Hutchinson, Gauss-Newton-Bartlett}\nSet m0 = 0, v0 = 0, h1−k = 0\nFor t = 1 to T do\n    Compute minibatch loss Lt(θt)\n    Compute gt = ∇Lt(θt)\n    mt = β1mt−1 + (1 − β1)gt\n    If t mod k = 1 then\n        Compute hˆt = Estimator(θt)\n        ht = β2ht−k + (1 − β2)hˆt\n    Else\n        ht = ht−1\n    θt = θt − ηtλθt (weight decay)\n    θt+1 = θt − ηt · clip(mt/ max{ht, ϵ}, ρ)\n\n```\n\n# Pytorch Implementation\n\n```python\n\nimport torch \n\nclass Sophia(torch.optim.Optimizer):\n    def __init__(self, model, input_data, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, k=10, estimator="Hutchinson", rho=1):\n        self.model = model\n        self.input_data = input_data\n        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, k=k, estimator=estimator, rho=rho)\n        super(Sophia, self).__init__(params, defaults)\n\n\n    def step(self, closure=None):\n        loss = None\n        if closure is not None:\n            loss = closure()\n\n        for group in self.param_groups:\n            for p in group["params"]:\n                if p.grad is None:\n                    continue\n                grad = p.grad.data\n                if grad.is_sparse:\n                    raise RuntimeError("Sophia does not support sparse gradients")\n                \n                state = self.state[p]\n\n                #state init\n                if len(state) == 0:\n                    state[\'step\'] = 0\n                    state[\'m\'] = torch.zeros_like(p.data)\n                    state[\'h\'] = torch.zeros_like(p.data)\n\n                m, h = state[\'m\'], state[\'h\']\n                beta1, beta2 = group[\'betas\']\n                state[\'step\'] += 1\n\n\n                if group[\'weight_decay\'] != 0:\n                    grad = grad.add(group["weight_decau"], p.data)\n\n                #update biased first moment estimate\n                m.mul_(beta1).add_(1 - beta1, grad)\n\n                #update hessian estimate\n                if state[\'step\'] % group[\'k\'] == 1:\n                    if group[\'estimator\'] == "Hutchinson":\n                        hessian_estimate = self.hutchinson(p, grad)\n                    elif group[\'estimator\'] == "Gauss-Newton-Bartlett":\n                        hessian_estimate = self.gauss_newton_bartlett(p, grad)\n                    else:\n                        raise ValueError("Invalid estimator choice")\n                    h.mul_(beta2).add_(1 - beta2, hessian_estimate)\n\n                #update params\n                p.data.add_(-group[\'lr\'] * group[\'weight_decay\'], p.data)\n                p.data.addcdiv_(-group[\'lr\'], m, h.add(group[\'eps\']).clamp(max=group[\'rho\']))\n\n        return loss\n    \n    def hutchinson(self, p, grad):\n        u = torch.randn_like(grad)\n        hessian_vector_product = torch.autograd.grad(grad.dot(u), p, retain_graph=True)[0]\n        return u * hessian_vector_product\n    \n    def gauss_newton_bartlett(self, p, grad):\n        B = len(self.input_data)\n        logits = [self.model(xb) for xb in self.input_data]\n        y_hats = [torch.softmax(logit, dim=0) for logit in logits]\n        g_hat = torch.autograd.grad(sum([self.loss_function(logit, y_hat) for logit, y_hat in zip(logits, y_hats)]) / B, p, retain_graph=True)[0]\n        return B * g_hat * g_hat\n    \n        \n```\n# Hyper-parameter Tuning Guide for Decoupled Sophia Optimizer\nWhen using the Decoupled Sophia optimizer, it\'s essential to tune the hyperparameters to achieve the best performance for your specific model and dataset. Here\'s a guide on how to tune the parameters for the Decoupled Sophia optimizer:\n\n## Learning Rate (lr)\nThe learning rate is a crucial hyperparameter that controls the step size of the parameter updates during the optimization process. In Decoupled Sophia, the update is written as p.data.addcdiv_(-group[\'lr\'], m, h.add(group[\'rho\'])), which is equivalent to the update in the paper up to a re-parameterization.\n\n### Tips for tuning the learning rate:\nChoose the learning rate to be about half the learning rate that you would use for AdamW. Some partial ongoing results indicate that the learning rate can be made even larger, possibly leading to faster convergence.\nRho (rho)\nThe rho parameter is used in the update rule to control the Hessian\'s influence on the parameter updates. It is essential to choose an appropriate value for rho to balance the trade-off between the gradient and the Hessian information.\n\n#### Tips for tuning rho:\nConsider choosing rho in the range of 0.03 to 0.04. The rho value seems transferable across different model sizes. For example, rho = 0.03 can be used in 125M and 335M Sophia-G models.\n\nThe (lr, rho) for 335M Sophia-G is chosen to be (2e-4, 0.03). Though we suspect that the learning rate can be larger, it\'s essential to experiment with different values to find the best combination for your specific use case.\n\n### Other Hyperparameters\nWhile the learning rate and rho are the most critical hyperparameters to tune, you may also experiment with other hyperparameters such as betas, weight_decay, and k (the frequency of Hessian updates). However, the default values provided in the optimizer should work well for most cases.\n\nRemember that hyperparameter tuning is an iterative process, and the best values may vary depending on the model architecture and dataset. Don\'t hesitate to experiment with different combinations and validate the performance on a held-out dataset or using cross-validation.\n\nFeel free to share your findings and experiences during hyperparameter tuning. Your valuable feedback and comments can help improve the optimizer and its usage in various scenarios.\n\n\n\n# Roadmap\nThe following roadmap outlines the future development plans for the Sophia optimizer. The roadmap is divided into three stages: short-term, mid-term, and long-term goals.\n\n\n## Short-term Goals\n\nReady to train plug in and play file with your own model or Andromeda\n\nPerformance improvements: Investigate and implement potential performance improvements to further reduce training time and computational resources -> Decoupled Sophia + heavy metric logging + Implement in Triton and or Jax?\n\nAdditional Hessian estimators: Research and implement other Hessian estimators to provide more options for users.\n\nHyperparameter tuning: Develop a set of recommended hyperparameters for various use cases and model architectures.\n\n# Mid-term Goals\nIntegration with Andromeda model: Train the Andromeda model using the Sophia optimizer and compare its performance with other optimizers.\n\nSophia optimizer variants: Explore and develop variants of the Sophia optimizer tailored for specific tasks, such as computer vision, multi-modality AI, and natural language processing, and reinforcement learning.\n\nDistributed training: Implement support for distributed training to enable users to train large-scale models using Sophia across multiple devices and nodes.\n\nAutomatic hyperparameter tuning: Develop an automatic hyperparameter tuning module to help users find the best hyperparameters for their specific use case.\n\n# Long-term Goals\nTraining multiple models in parallel: Develop a framework for training multiple models concurrently with different optimizers, allowing users to test and compare the performance of various optimizers, including Sophia, on their specific tasks.\n\nSophia optimizer for other domains: Adapt the Sophia optimizer for other domains, such as optimization in reinforcement learning, Bayesian optimization, and evolutionary algorithms.\n\n\nBy following this roadmap, we aim to make the Sophia optimizer a powerful and versatile tool for the deep learning community, enabling users to train their models more efficiently and effectively.\n\n\n# Epoch 1 - Decoupled Sophia\n\nThe DecoupledSophia optimizer offers several benefits that can help accelerate training speed and improve the flexibility of the optimization process:\n\nModularity: Decoupling the Hessian estimation from the main optimizer allows users to easily plug in different Hessian estimators without modifying the core optimizer code. This modularity makes it easier to experiment with various Hessian estimation techniques and find the best one for a specific task.\n\nEase of experimentation: With a decoupled architecture, researchers and practitioners can develop and test new Hessian estimators independently from the optimizer. This can lead to faster innovation and the discovery of more efficient Hessian estimation methods, which can further accelerate training speed.\n\nCustomization: Users can create custom Hessian estimators tailored to their specific use cases or model architectures. This customization can potentially lead to better optimization performance and faster training times.\n\nImproved maintainability: Separating the Hessian estimation from the optimizer makes the codebase easier to maintain and understand. This can lead to faster bug fixes and improvements in the optimizer\'s performance.\n\nBy offering these benefits, DecoupledSophia can help users accelerate training speed and improve the overall optimization process. The modular design allows for easy experimentation with different Hessian estimators, which can lead to the discovery of more efficient techniques and ultimately faster training times.\n',
    'author': 'Kye Gomez',
    'author_email': 'kye@apac.ai',
    'maintainer': 'None',
    'maintainer_email': 'None',
    'url': 'https://github.com/kyegomez/Sophia',
    'packages': packages,
    'package_data': package_data,
    'install_requires': install_requires,
    'python_requires': '>=3.6,<4.0',
}


setup(**setup_kwargs)
