{"id":13456598,"url":"https://github.com/LargeWorldModel/LWM","last_synced_at":"2025-03-24T11:30:47.813Z","repository":{"id":222416898,"uuid":"754447605","full_name":"LargeWorldModel/LWM","owner":"LargeWorldModel","description":null,"archived":false,"fork":false,"pushed_at":"2024-10-15T16:01:16.000Z","size":4200,"stargazers_count":7109,"open_issues_count":60,"forks_count":550,"subscribers_count":66,"default_branch":"main","last_synced_at":"2024-10-16T11:41:09.703Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":null,"language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/LargeWorldModel.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-08T04:16:42.000Z","updated_at":"2024-10-16T07:30:47.000Z","dependencies_parsed_at":"2024-07-13T20:01:09.340Z","dependency_job_id":"8c81766d-db76-45e0-aa53-3ad014821233","html_url":"https://github.com/LargeWorldModel/LWM","commit_stats":null,"previous_names":["largeworldmodel/lwm"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/LargeWorldModel%2FLWM","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/LargeWorldModel%2FLWM/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/LargeWorldModel%2FLWM/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/LargeWorldModel%2FLWM/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/LargeWorldModel","download_url":"https://codeload.github.com/LargeWorldModel/LWM/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":221962414,"owners_count":16908336,"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","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":[],"created_at":"2024-07-31T08:01:24.697Z","updated_at":"2024-10-29T00:30:42.539Z","avatar_url":"https://github.com/LargeWorldModel.png","language":"Python","funding_links":[],"categories":["Python","Project List","HarmonyOS","多模态大模型"],"sub_categories":["\u003cspan id=\"tool\"\u003eLLM (LLM \u0026 Tool)\u003c/span\u003e","Windows Manager","网络服务_其他"],"readme":"# Large World Model (LWM)\n\n[[Project]](https://largeworldmodel.github.io/)\n[[Paper]](https://arxiv.org/abs/2402.08268)\n[[Models]](https://huggingface.co/LargeWorldModel)\n\n**Large World Model (LWM)** is a general-purpose large-context multimodal autoregressive model. It is trained on a large dataset of diverse long videos and books using RingAttention, and can perform language, image, and video understanding and generation.\n\n\n## Approach\n\n\u003cdiv align=\"center\"\u003e\n  \u003cimg src=\"./imgs/data.png\"/\u003e\n\u003c/div\u003e\n\nCurrent language models fall short in understanding aspects of the world not easily described in words, and struggle with complex, long-form tasks. Video sequences offer valuable temporal information absent in language and static images, making them attractive for joint modeling with language. Such models could develop a understanding of both human textual knowledge and the physical world, enabling broader AI capabilities for assisting humans. However, learning from millions of tokens of video and language sequences poses challenges due to memory constraints, computational complexity, and limited datasets. To address these challenges, we curate a large dataset of diverse videos and books, utilize the RingAttention technique to scalably train on long sequences, and gradually increase context size from 4K to 1M tokens. This paper makes the following contributions: (a) Largest context size neural network: We train one of the largest context size transformers on long video and language sequences, setting new benchmarks in difficult retrieval tasks and long video understanding. (b) Solutions for overcoming vision-language training challenges, including using masked sequence packing for mixing different sequence lengths, loss weighting to balance language and vision, and model-generated QA dataset for long sequence chat. (c) A highly-optimized implementation with RingAttention, masked sequence packing, and other key features for training on millions-length multimodal sequences. (d) Fully open-sourced a family of 7B parameter models capable of processing long text documents (LWM-Text, LWM-Text-Chat) and videos (LWM, LWM-Chat) of over 1M tokens.\nThis work paves the way for training on massive datasets of long video and language to develop understanding of both human knowledge and the multimodal world, and broader capabilities.\n\n## LWM Capabilities\n\n\u003cdiv align=\"center\"\u003e\n  \u003cimg src=\"./imgs/single_needle_1M.png\"/\u003e\n  \u003cp\u003e\n  LWM can retrieval facts across 1M context with high accuracy.\n  \u003c/p\u003e\n\u003c/div\u003e\n\n\u003cbr /\u003e\n\n\u003cdiv align=\"center\"\u003e\n  \u003cimg src=\"./imgs/long_video_chat_main.png\"/\u003e\n  \u003cp\u003e\n  LWM can answer questions over 1 hour YouTube video.\n  \u003c/p\u003e\n\u003c/div\u003e\n\n\u003cbr /\u003e\n\n\u003cdiv align=\"center\"\u003e\n  \u003cimg src=\"./imgs/image_chat.png\"/\u003e\n  \u003cp\u003e\n  LWM can chat with images.\n  \u003c/p\u003e\n\u003c/div\u003e\n\n\u003cbr /\u003e\n\n\u003cdiv align=\"center\"\u003e\n  \u003cimg src=\"./imgs/image_video_gen.png\"/\u003e\n  \u003cp\u003e\n  LWM can generate videos and images from text.\n  \u003c/p\u003e\n\u003c/div\u003e\n\n\n## Setup\n\nThis codebase is supported on Ubuntu and has not been tested on Windows or macOS. We recommend using TPUs for training and inference, although it is also possible to use GPUs. On TPU, the code is highly optimized with Jax's Pallas and can achieve high MFUs with RingAttention at very large context sizes. On GPU, the code is based on XLA and is not as optimized as it is for TPU.\n\nInstall the requirements with:\n```\nconda create -n lwm python=3.10\nconda activate lwm\npip install -r gpu_requirements.txt\n```\nor set up TPU VM with:\n```\nsh tpu_requirements.sh\n```\n\n\n## Available models\n\nThere are language-only and video-language versions, offering context sizes from 32K, to 128K, 256K and 1M tokens. The vision-language models are available only in Jax, and the language-only models are available in both PyTorch and Jax. Below are the names of the available models and their corresponding context sizes and capabilities:\n\n| Model Name         | Context Size | Language or Vision-Language | Chat or Base | URL                                                                                                                                          |\n|--------------------|--------------|-----------------------------|--------------|----------------------------------------------------------------------------------------------------------------------------------------------|\n| LWM-Text-Chat-128K | 128K         | Language                    | Chat         | [[Pytorch](https://huggingface.co/LargeWorldModel/LWM-Text-Chat-128K)][[Jax](https://huggingface.co/LargeWorldModel/LWM-Text-Chat-128K-Jax)] |\n| LWM-Text-Chat-256K | 256K         | Language                    | Chat         | [[Pytorch](https://huggingface.co/LargeWorldModel/LWM-Text-Chat-256K)][[Jax](https://huggingface.co/LargeWorldModel/LWM-Text-Chat-256K-Jax)] |\n| LWM-Text-Chat-512K | 512K         | Language                    | Chat         | [[Pytorch](https://huggingface.co/LargeWorldModel/LWM-Text-Chat-512K)][[Jax](https://huggingface.co/LargeWorldModel/LWM-Text-Chat-512K-Jax)] |\n| LWM-Text-Chat-1M   | 1M           | Language                    | Chat         | [[Pytorch](https://huggingface.co/LargeWorldModel/LWM-Text-Chat-1M)][[Jax](https://huggingface.co/LargeWorldModel/LWM-Text-Chat-1M-Jax)]     |\n| LWM-Text-128K      | 128K         | Language                    | Base         | [[Pytorch](https://huggingface.co/LargeWorldModel/LWM-Text-128K)][[Jax](https://huggingface.co/LargeWorldModel/LWM-Text-128K-Jax)]           |\n| LWM-Text-256K      | 256K         | Language                    | Base         | [[Pytorch](https://huggingface.co/LargeWorldModel/LWM-Text-256K)][[Jax](https://huggingface.co/LargeWorldModel/LWM-Text-256K-Jax)]           |\n| LWM-Text-512K      | 512K         | Language                    | Base         | [[Pytorch](https://huggingface.co/LargeWorldModel/LWM-Text-512K)][[Jax](https://huggingface.co/LargeWorldModel/LWM-Text-512K-Jax)]           |\n| LWM-Text-1M        | 1M           | Language                    | Base         | [[Pytorch](https://huggingface.co/LargeWorldModel/LWM-Text-1M)][[Jax](https://huggingface.co/LargeWorldModel/LWM-Text-1M-Jax)]               |\n| LWM-Chat-32K       | 32K          | Vision-Language             | Chat         | [[Jax](https://huggingface.co/LargeWorldModel/LWM-32K-Jax)]                                                                                  |\n| LWM-Chat-128K      | 128K         | Vision-Language             | Chat         | [[Jax](https://huggingface.co/LargeWorldModel/LWM-128K-Jax)]                                                                                 |\n| LWM-Chat-1M        | 1M           | Vision-Language             | Chat         | [[Jax](https://huggingface.co/LargeWorldModel/LWM-1M-Jax)]                                                                                   |\n\n\n## Code structure\nUse `scan_query_chunk_size` and `scan_key_chunk_size` to control the block size in blockwise compute of the self-attention. Use `scan_mlp_chunk_size` to control the block size in blockwise compute of the feedforward network. Use `scan_attention=True` and `scan_mlp=True` to enable/disable blockwise compute in the self-attention and feed-forward network.\n\nYou can use `mesh_dim=dp, fsdp, tp, sp` to control the degree of parallelism and RingAttention. It is a string of 4 integers separated by commas, representing the number of data parallelism, fully sharded data parallelism, tensor parallelism, and sequence parallelism.\nFor example, `mesh_dim='1,64,4,1'` means 1 data parallelism, 64 fully sharded data parallelism, 4 tensor parallelism, and 1 sequence parallelism. `mesh_dim='1,1,4,64'` means 1 data parallelism, 1 fully sharded data parallelism, 4 tensor parallelism, and 64 sequence parallelism for RingAttention.\n\n\n## Running Jax Models\nIn this section, we provide instructions on how to run each of the provided scripts. For each script, you may need to fill in your own paths and values in the variables described in the beginning of each script.\n\nTo run each of the following scripts, use `bash \u003cscript_name\u003e.sh`:\n- Language model training: `bash scripts/run_train_text.sh`\n- Vision-Language model training: `bash scripts/run_train_vision_text.sh`\n- Single Needle Evals (Language Model): `bash scripts/run_eval_needle.sh`\n- Multi Needle Evals (Language Model): `bash scripts/run_eval_needle_multi.sh`\n- Sampling images (Vision-Language Model): `bash scripts/run_sample_image.sh`\n- Sampling videos (Vision-LanguageModel): `bash scripts/run_sample_video.sh`\n- Image / Video understanding (Vision-Language Model): `bash scripts/run_vision_chat.sh`\n\nBy default the `mesh_dim` argument puts all devices on `tp` (tensor parallelism). For longer sequences, you may want to include `sp`, which is the last dimension in the `mesh_dim`.\n\nWhen running needle evals, you may need to adjust the `theta` and `max_sequence_length` arguments in the scripts depending on the model. Below shows the correct values for each model.\n\n|                     | LWM-Text-128K /  LWM-Text-Chat-128K | LWM-Text-256K /  LWM-Text-Chat-256K | LWM-Text-512K / LWM-Text-Chat-512K | LWM-Text-1M / LWM-Text-Chat-1M |\n|---------------------|:-----------------------------------:|:-----------------------------------:|:----------------------------------:|:------------------------------:|\n| theta               |               10000000              |               10000000              |              25000000              |            50000000            |\n| max_sequence_length |                131072               |                262144               |               524288               |             1048576            |\n\n\nAn example of filling out a script (`run_sample_video.sh`) is as follows\n```bash\n#! /bin/bash\n\nexport SCRIPT_DIR=\"$( cd -- \"$( dirname -- \"${BASH_SOURCE[0]}\" )\" \u0026\u003e /dev/null \u0026\u0026 pwd )\"\nexport PROJECT_DIR=\"$( cd -- \"$( dirname -- \"$SCRIPT_DIR\" )\" \u0026\u003e /dev/null \u0026\u0026 pwd )\"\ncd $PROJECT_DIR\nexport PYTHONPATH=\"$PYTHONPATH:$PROJECT_DIR\"\n\nexport llama_tokenizer_path=\"LargeWorldModel/LWM-Text-1M\"\nexport vqgan_checkpoint=\"/path/to/ckpt/folder/vqgan\"\nexport lwm_checkpoint=\"params::/path/to/ckpt/folder/params\"\n\npython3 -u -m lwm.vision_generation \\\n    --prompt='Fireworks over the city' \\\n    --output_file='fireworks.mp4' \\\n    --temperature_image=1.0 \\\n    --temperature_video=1.0 \\\n    --top_k_image=8192 \\\n    --top_k_video=1000 \\\n    --cfg_scale_image=5.0 \\\n    --cfg_scale_video=1.0 \\\n    --vqgan_checkpoint=\"$vqgan_checkpoint\" \\\n    --n_frames=8 \\\n    --mesh_dim='!1,1,-1,1' \\\n    --dtype='fp32' \\\n    --load_llama_config='7b' \\\n    --update_llama_config=\"dict(sample_mode='vision',theta=50000000,max_sequence_length=32768,scan_attention=False,scan_query_chunk_size=128,scan_key_chunk_size=128,scan_mlp=False,scan_mlp_chunk_size=8192,scan_layers=True)\" \\\n    --load_checkpoint=\"$lwm_checkpoint\" \\\n    --tokenizer=\"$llama_tokenizer_path\"\nread\n```\n\n\n## Needle Haystack Data\nRun `python scripts/create_needle_data.py`\n\n\n## Running PyTorch Models\nOnly text and text chat models are currently supported for PyTorch inference. PyTorch models can be loaded as Hugging Face `LlamaForCausalLM` models. Run `python scripts/sample_pyt.py` to sample. You may need to separately install `torch`.\n\n## Documentation\n\nFor more details on the codebase, please refer to the [data.md](docs/data.md) and [sharding.md](docs/sharding.md).\nThe [data.md](docs/data.md) provides details on the data processing and the [sharding.md](docs/sharding.md) provides details on the sharding and parallelism.\n\n\n## If you have issues\n\nThis is based on the [codebase](https://github.com/haoliuhl/ringattention) of RingAttention, with the necessary features for vision-language training. The training and inference have been tested on both TPUv3 and TPUv4.\n\nIf you encounter bugs, please open a GitHub issue!\n\n\n## Citation\n\nIf you use this codebase, or otherwise found our work valuable, please cite:\n\n```\n@article{liu2023world,\n    title={World Model on Million-Length Video and Language with RingAttention},\n    author={Liu, Hao and Yan, Wilson and Zaharia, Matei and Abbeel, Pieter},\n    journal={arXiv preprint},\n    year={2024},\n}\n@article{liu2023ring,\n    title={Ring Attention with Blockwise Transformers for Near-Infinite Context},\n    author={Liu, Hao and Zaharia, Matei and Abbeel, Pieter},\n    journal={International Conference on Learning Representations},\n    year={2024}\n}\n@article{liu2023blockwise,\n    title={Blockwise Parallel Transformer for Large Context Models},\n    author={Liu, Hao and Abbeel, Pieter},\n    journal={Advances in neural information processing systems},\n    year={2023}\n}\n```\n\n## License\n\nLWM's code is released under the Apache 2.0 License. See [LICENSE](https://github.com/LargeWorldModel/lwm/blob/main/LICENSE) for further details. The models are released under the Llama-2 license.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2FLargeWorldModel%2FLWM","html_url":"https://awesome.ecosyste.ms/projects/github.com%2FLargeWorldModel%2FLWM","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2FLargeWorldModel%2FLWM/lists"}