{"id":18401676,"url":"https://github.com/borealisai/flora-opt","last_synced_at":"2025-09-02T14:30:46.011Z","repository":{"id":220952322,"uuid":"752948912","full_name":"BorealisAI/flora-opt","owner":"BorealisAI","description":"This is the official repository for the paper \"Flora: Low-Rank Adapters Are Secretly Gradient Compressors\" in ICML 2024.","archived":false,"fork":false,"pushed_at":"2024-07-01T21:07:59.000Z","size":688,"stargazers_count":104,"open_issues_count":0,"forks_count":5,"subscribers_count":2,"default_branch":"main","last_synced_at":"2025-08-16T12:53:31.033Z","etag":null,"topics":["deep-learning","flax","jax","large-language-models","lora","memory-efficient-tuning","optax","peft","random-projection","transformers"],"latest_commit_sha":null,"homepage":"https://arxiv.org/abs/2402.03293","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"lgpl-3.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/BorealisAI.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2024-02-05T07:01:09.000Z","updated_at":"2025-06-14T07:37:48.000Z","dependencies_parsed_at":"2024-06-20T03:57:36.116Z","dependency_job_id":"aa0c0941-bc02-4cfc-8d2b-0f59bab80875","html_url":"https://github.com/BorealisAI/flora-opt","commit_stats":null,"previous_names":["manga-uofa/flora","borealisai/flora-opt"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/BorealisAI/flora-opt","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BorealisAI%2Fflora-opt","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BorealisAI%2Fflora-opt/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BorealisAI%2Fflora-opt/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BorealisAI%2Fflora-opt/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/BorealisAI","download_url":"https://codeload.github.com/BorealisAI/flora-opt/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BorealisAI%2Fflora-opt/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":273297477,"owners_count":25080405,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","status":"online","status_checked_at":"2025-09-02T02:00:09.530Z","response_time":77,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["deep-learning","flax","jax","large-language-models","lora","memory-efficient-tuning","optax","peft","random-projection","transformers"],"created_at":"2024-11-06T02:39:37.140Z","updated_at":"2025-09-02T14:30:45.666Z","avatar_url":"https://github.com/BorealisAI.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# 💐 Flora: Low-Rank Adapters Are Secretly Gradient Compressors\n\nThis is the official repository for the paper [Flora: Low-Rank Adapters Are Secretly Gradient Compressors](https://arxiv.org/abs/2402.03293) in ICML 2024 ([research blog post on Flora](https://www.borealisai.com/research-blogs/pre-training-multi-billion-parameter-llms-on-a-single-gpu-with-flora/)).  This repository contains the code for the experiments in the paper.\n\nFlora dramatically decreases the GPU memory needed to pre-train and fine-tune models without compromising quality or causing significant slowdown. \n\n![Flora Graphs](https://github.com/BorealisAI/flora-opt/assets/1929273/5bb88165-9bb6-4dbb-b346-5f172356d434)\n\n## Installation\n\nYou can install the library using pip:\n\n```bash\npip install 'flora-opt[torch]' # for PyTorch\n```\n\nor\n\n```bash\npip install 'flora-opt[jax]' # for JAX\n```\n\n## Usage\n\nThe library is designed to be compatible with huggingface libraries. You can use it as a drop-in replacement for huggingface libraries. Here is an example for PyTorch:\n\n```diff\n- optimizer = transformers.AdamW(model.parameters(), lr=1e-5)\n+ optimizer = flora_opt.Flora(model.parameters(), lr=1e-5)\n- accelerator = accelerate.Accelerator(**kwargs)\n+ accelerator = flora_opt.FloraAccelerator(**kwargs)\n```\n\nEverything else remains the same. You can find more examples in the `examples` folder.\n\n## How it works\n\nNormally, there are three components in deep learning training: the model parameters, the optimizer, and the activations. Take the Adam optimizer as an example, the overall procedure is as follows:\n\n| Procedure                                                | Memory                                                   |\n| -------------------------------------------------------- | -------------------------------------------------------- |\n| \u003cimg src=\"assets/figures/Adam-Steps.gif\" width=\"100%\" /\u003e | \u003cimg src=\"assets/figures/Adam-Memory.gif\" width=\"81%\" /\u003e |\n\nSince Adam needs to store the first- and second-order moments for each parameter, the memory may double the size of the model.\n\nIn our work, instead of maintaining the optimizers' states, we propose to use low-rank random projections to compress the moments. The overall procedure is as follows:\n\n| Procedure                                                 | Memory                                                    |\n| --------------------------------------------------------- | --------------------------------------------------------- |\n| \u003cimg src=\"assets/figures/Flora-Steps.gif\" width=\"100%\" /\u003e | \u003cimg src=\"assets/figures/Flora-Memory.gif\" width=\"81%\" /\u003e |\n\nThe low-rank random projections reduces the memory usage by a factor of `rank / d`, where `rank` is the rank of the low-rank random projections and `d` is the dimension of the model parameters. In addition, the low-rank random projections can be generated on-the-fly using the random seed, which further reduces the memory usage.\n\nMoreover, Flora is compatible with the existing memory-efficient training techniques. For example, Flora can be combined with the activation checkpointing (AC) and layer-by-layer update (LOMO) to further reduce the memory usage, as shown below:\n\n| Procedure                                                         | Memory                                                            |\n| ----------------------------------------------------------------- | ----------------------------------------------------------------- |\n| \u003cimg src=\"assets/figures/Flora-AC-LOMO-Steps.gif\" width=\"100%\" /\u003e | \u003cimg src=\"assets/figures/Flora-AC-LOMO-Memory.gif\" width=\"81%\" /\u003e |\n\nIn summary, Flora is a simple yet effective method to compress the optimizer's states, which can be easily integrated into existing training frameworks.\n\n## Paper Replications (with JAX)\n\nTo replicate major experiments in the paper, run the following commands:\n\n```bash\npip install -r examples/flax/requirements.txt\nsh replicate.sh\n```\n\nYou can also run individual experiments by selecting the corresponding script in the file `replicate.sh`.\n\n## Explanation\n\nThe arguments for the `Flora` optimizer (for PyTorch) are explained below:\n\n```python\nflora_opt.Flora(\n    params: Union[Iterable[torch.Tensor], Iterable[Dict[str, Any]]],  # model parameters\n    lr: float = None,  # learning rate\n    rank: int = None,  # rank of the low-rank random projections\n    kappa: int = 1000,  # the interval for updating the low-rank random projections\n    eps: tuple[float, float] = (1e-30, 1e-3),  # Adafactor parameter\n    clip_threshold: float = 1.0,  # Adafactor parameter\n    decay_rate: float = -0.8,  # decay rate in Adafactor\n    beta1: Optional[float] = None,  # decay rate for the first moment\n    weight_decay: float = 0.0,  # weight decay coefficient\n    scale_parameter: bool = True,  # Adafactor parameter\n    relative_step: bool = False,  # Adafactor parameter\n    warmup_init: bool = False,  # Adafactor parameter\n    factorize_second_moment: bool = True,  # use Adafactor or Adam\n    seed: int = 0,  # random seed to generate the low-rank random projections\n    quantization: bool = False,  # whether to quantize the states\n)\n```\n\nFor JAX, the arguments are similar. The translation can be found in `flora_opt/optimizers/torch/__init__.py`.\n\n### Citation\n\n```bibtex\n@inproceedings{hao2024flora,\n  title={Flora: Low-Rank Adapters Are Secretly Gradient Compressors},\n  author={Hao, Yongchang and Cao, Yanshuai and Mou, Lili},\n  booktitle={Forty-first International Conference on Machine Learning},\n  url={https://arxiv.org/abs/2402.03293},\n  year={2024}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fborealisai%2Fflora-opt","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fborealisai%2Fflora-opt","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fborealisai%2Fflora-opt/lists"}