Update README.md (EasyDeL auto-generated)

README.md

@@ -1,115 +1,167 @@
 ---
+library_name: easydel
+pipeline_tag: text-generation
 tags:
-- EasyDeL
-- KimiVLForConditionalGeneration
-- TaskType.IMAGE_TEXT_TO_TEXT
-- AttentionMechanisms.RAGGED_PAGE_ATTENTION_V3
-- safetensors
-- TPU
-- GPU
-- XLA
-- Flax
+  - easydel
+  - jax
+  - "kimi_vl"
+  - "CausalLM"
+  - "auto"
 ---
+
 <p align="center">
   <a href="https://github.com/erfanzar/EasyDeL">
-    <img src="https://raw.githubusercontent.com/erfanzar/easydel/main/images/easydel-logo-with-text.png" height="80">
+    <img src="https://raw.githubusercontent.com/erfanzar/easydel/main/images/easydel-logo-with-text.png" height="80" alt="EasyDeL" />
   </a>
 </p>
 
+<h1 align="center">Kimi-VL-A3B-Instruct</h1>
+
 <p align="center">
-  <a href="https://github.com/erfanzar/EasyDeL">
-    <img src="https://img.shields.io/badge/🤗_EasyDeL-0.2.0-blue.svg" />
+  A model compatible with the EasyDeL JAX stack.
+</p>
+
+<p align="center">
+  <a href="https://huggingface.co/EasyDeL/Kimi-VL-A3B-Instruct">
+    <img src="https://img.shields.io/static/v1?label=HF&message=EasyDeL/Kimi-VL-A3B-Instruct&color=FFD21E&style=flat-square" alt="HuggingFace Repo" />
   </a>
   <a href="https://github.com/erfanzar/EasyDeL">
-    <img src="https://img.shields.io/badge/Model_Type-KimiVLForConditionalGeneration-green.svg" />
+    <img src="https://img.shields.io/static/v1?label=EasyDeL&message=v0.2.0&color=0B5FFF&style=flat-square" alt="EasyDeL Version" />
   </a>
+  <img src="https://img.shields.io/static/v1?label=Model&message=kimi_vl&color=0A66C2&style=flat-square" alt="Model Type" />
+  <img src="https://img.shields.io/static/v1?label=Task&message=CausalLM&color=2EAD4D&style=flat-square" alt="Task" />
+  <img src="https://img.shields.io/static/v1?label=Attention&message=auto&color=8A2BE2&style=flat-square" alt="Attention Mechanism" />
 </p>
 
-# EasyDeL/Kimi-VL-A3B-Instruct
-
-A model implemented using the EasyDeL framework, designed to deliver optimal performance for large-scale natural language processing tasks.
+---
 
 ## Overview
 
-This model is built using [EasyDeL](https://github.com/erfanzar/EasyDeL), an open-source framework designed to enhance and streamline the training and serving process of machine learning models, with a primary focus on Jax/Flax on TPU/GPU at scale.
+This checkpoint is intended to be loaded with EasyDeL on JAX (CPU/GPU/TPU). It supports sharded loading with `auto_shard_model=True` and configurable precision via `dtype`, `param_dtype`, and `precision`.
 
-EasyDeL provides an efficient, highly-optimized, and customizable machine learning model compatible with both GPU and TPU environments. Built with JAX, this model supports advanced features such as sharded model parallelism, making it suitable for distributed training and inference and customized kernels.
+## Quickstart
 
-## Features Provided by EasyDeL
+```python
+import easydel as ed
+from jax import numpy as jnp, lax
 
-**EasyDeL Framework Features:**
+repo_id = "EasyDeL/Kimi-VL-A3B-Instruct"
+
+dtype = jnp.bfloat16  # try jnp.float16 on many GPUs
+
+model = ed.AutoEasyDeLModelForCausalLM.from_pretrained(
+    repo_id,
+    dtype=dtype,
+    param_dtype=dtype,
+    precision=lax.Precision("fastest"),
+    config_kwargs=ed.EasyDeLBaseConfigDict(
+        attn_dtype=dtype,
+        attn_mechanism=ed.AttentionMechanisms.AUTO,
+        fsdp_is_ep_bound=True,
+        sp_is_ep_bound=True,
+        moe_method=ed.MoEMethods.FUSED_MOE,
+    ),
+    auto_shard_model=True,
+    partition_axis=ed.PartitionAxis(),
+)
+```
 
-- **Efficient Implementation**: Built with JAX/Flax for high-performance computation.
-- **Modern Architecture**: Built on Flax NNX for better integration, modularity, and performance.
-- **Multi-Device Support**: Optimized to run on TPU, GPU, and CPU environments.
-- **Sharded Model Parallelism**: Supports model parallelism across multiple devices for scalability (using `auto_shard_model=True`).
-- **Customizable Precision**: Allows specification of `dtype`, `param_dtype`, and `precision`.
-- **Advanced Serving**: Includes `eSurge` LLM serving engine, `vWhisper` speech endpoints, and OpenAI-compatible APIs.
-- **Optimized Kernels**: Integrates multiple attention mechanisms (like `AttentionMechanisms.RAGGED_PAGE_ATTENTION_V3`) and platform-specific optimizations.
+If the repository only provides PyTorch weights, pass `from_torch=True` to `from_pretrained(...)`.
 
-## Installation
+## Sharding & Parallelism (Multi-Device)
 
-To use this model via EasyDeL, first install EasyDeL:
+EasyDeL can scale to multiple devices by creating a logical device mesh. Most EasyDeL loaders use a 5D mesh:
 
-```bash
-pip install easydel
-```
+- `dp`: data parallel (replicated parameters, different batch shards)
+- `fsdp`: parameter sharding (memory saver; often the biggest axis)
+- `ep`: expert parallel (MoE; keep `1` for non-MoE models)
+- `tp`: tensor parallel (splits large matmuls)
+- `sp`: sequence parallel (splits sequence dimension)
+
+Use `sharding_axis_names=("dp","fsdp","ep","tp","sp")` and choose `sharding_axis_dims` so that their product equals your device count.
+You can use `-1` in `sharding_axis_dims` to let EasyDeL infer the remaining dimension.
 
-## Usage
+<details>
+<summary>Example sharding configs</summary>
+
+```python
+# 8 devices, pure FSDP
+sharding_axis_dims = (1, 8, 1, 1, 1)
+
+# 8 devices, 2-way DP x 4-way FSDP
+sharding_axis_dims = (2, 4, 1, 1, 1)
+
+# 8 devices, 4-way FSDP x 2-way TP
+sharding_axis_dims = (1, 4, 1, 2, 1)
+```
+</details>
 
-### Loading the Pre-trained Model
+## Using via `eLargeModel` (ELM)
 
-To load this pre-trained model with EasyDeL:
+`eLargeModel` is a higher-level interface that wires together loading, sharding, training, and eSurge inference from a single config.
 
 ```python
-from easydel import AutoEasyDeLModelForCausalLM, EasyDeLBaseConfigDict, AttentionMechanisms
-from jax import numpy as jnp, lax
+from easydel import eLargeModel
 
-# Define max_length if needed for memory optimization
-max_length = None
-
-# Load model and parameters
-# Set auto_shard_model=True to automatically distribute across devices
-model = AutoEasyDeLModelForCausalLM.from_pretrained(
-    "EasyDeL/Kimi-VL-A3B-Instruct",
-    config_kwargs=EasyDeLBaseConfigDict(
-        # use_scan_mlp=False, # Set to True to potentially reduce memory usage
-        attn_dtype=jnp.float16, # Or jnp.bfloat16
-        # freq_max_position_embeddings=max_length, # Set if using RoPE and need truncation
-        # mask_max_position_embeddings=max_length, # Set if max length is defined
-        attn_mechanism=AttentionMechanisms.RAGGED_PAGE_ATTENTION_V3 # Matches the mechanism used by this model
-    ),
-    dtype=jnp.float16, # Or jnp.bfloat16 - Computation data type
-    param_dtype=jnp.float16, # Or jnp.bfloat16 - Parameter data type
-    precision=lax.Precision("fastest"), # Like "default", "fastest", "high", "highest"
-    auto_shard_model=True, # Auto-shard across available devices
-)
+repo_id = "EasyDeL/Kimi-VL-A3B-Instruct"
+
+elm = eLargeModel.from_pretrained(repo_id)  # task is auto-detected
+elm.set_dtype("bf16")
+elm.set_sharding(axis_names=("dp", "fsdp", "ep", "tp", "sp"), axis_dims=(1, -1, 1, 1, 1))
+
+model = elm.build_model()
+# Optional: build an inference engine
+# engine = elm.build_esurge()
 ```
 
-## Supported Tasks
+<details>
+<summary>ELM YAML config example</summary>
 
-The primary task for this model is **TaskType.IMAGE_TEXT_TO_TEXT**. Further specific supported tasks are not explicitly listed.
+```yaml
+model:
+  name_or_path: "EasyDeL/Kimi-VL-A3B-Instruct"
 
-## Limitations
+loader:
+  dtype: bf16
+  param_dtype: bf16
 
-**General Limitations:**
+sharding:
+  axis_dims: [1, -1, 1, 1, 1]
+  auto_shard_model: true
+```
+</details>
 
-- **Hardware Dependency**: Performance can vary significantly based on the hardware (TPU/GPU) used.
-- **JAX/Flax Setup Required**: The environment must support JAX/Flax for optimal use.
-- **Experimental Features**: Some EasyDeL features (like custom kernels) may require additional configuration.
+## Features
 
-## License 📜
+**EasyDeL:**
+- JAX native implementation and sharded execution
+- Configurable attention backends via `AttentionMechanisms.*`
+- Precision control via `dtype`, `param_dtype`, and `precision`
 
-EasyDeL is released under the Apache v2 license. The license for this specific model might differ; please consult the original model repository or documentation.
+## Installation
 
-```code
-# Apache License 2.0 (referring to EasyDeL Framework)
-# ... (Full license text usually included in the main repo) ...
+```bash
+pip install easydel
 ```
 
-## Citation
+## Links
 
-If you use EasyDeL in your research or work, please cite it:
+- EasyDeL GitHub: https://github.com/erfanzar/EasyDeL
+- Docs: https://easydel.readthedocs.io/en/latest/
+
+## Supported Tasks
+
+- CausalLM
+
+## Limitations
+
+- Refer to the original model card for training data, evaluation, and intended use.
+
+## License
+
+EasyDeL is released under the Apache-2.0 license. The license for this model's weights may differ; please consult the original repository.
+
+## Citation
 
 ```bibtex
 @misc{Zare Chavoshi_2023,
@@ -119,5 +171,3 @@ If you use EasyDeL in your research or work, please cite it:
     year={2023}
 }
 ```
-
-Please also consider citing the original paper or source for the **EasyDeL/Kimi-VL-A3B-Instruct** model architecture if applicable.