From a3023cb4a0cc537984b528739122495b25dfe4fe Mon Sep 17 00:00:00 2001 From: Jen Ha Date: Tue, 18 Nov 2025 14:33:04 +0000 Subject: [PATCH] remove hw from tables for clarity --- bonsai/models/densenet121/README.md | 9 +++---- bonsai/models/efficientnet/README.md | 14 ++++------- bonsai/models/llada_8b/README.md | 11 ++++----- bonsai/models/qwen3/README.md | 36 ++++++++++++---------------- bonsai/models/resnet/README.md | 8 +++---- bonsai/models/sam2/README.md | 18 +++++++------- bonsai/models/unet/README.md | 13 ++++------ bonsai/models/vae/README.md | 13 ++++------ bonsai/models/vgg19/README.md | 9 +++---- bonsai/models/vit/README.md | 8 +++---- 10 files changed, 62 insertions(+), 77 deletions(-) diff --git a/bonsai/models/densenet121/README.md b/bonsai/models/densenet121/README.md index a3ad7f98..fd9dd781 100644 --- a/bonsai/models/densenet121/README.md +++ b/bonsai/models/densenet121/README.md @@ -2,10 +2,11 @@ This directory contains a pure JAX implementation of the [DenseNet-121](https://huggingface.co/keras/densenet_121_imagenet) model, using the Flax NNX API. -## Tested on -| Model Name | Config | CPU | GPU A100 (1x) | GPU H100 (1x) | GPU A100 (8x) | GPU H100 (8x) | TPU v2 (8x) | TPU v5e (1x) | -|:--------------------------------------------------------| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | -| [DenseNet-121](https://huggingface.co/keras/densenet_121_imagenet) | ✅ Supported | ✅ Runs | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | +## Model Configuration Support Status + +| Model Name | Config Support Status | +| :--- | :--- | +| [DenseNet-121](https://huggingface.co/keras/densenet_121_imagenet) | **✅ Supported** | ### Running this model diff --git a/bonsai/models/efficientnet/README.md b/bonsai/models/efficientnet/README.md index 679dbb16..f09a6884 100644 --- a/bonsai/models/efficientnet/README.md +++ b/bonsai/models/efficientnet/README.md @@ -3,16 +3,12 @@ This directory contains a pure JAX implementation of the [Efficientnet](https://arxiv.org/abs/1905.11946), using the [Flax NNX](https://flax.readthedocs.io/en/v0.8.3/experimental/nnx/index.html) API. -## Tested on: -*(Last Updated: 2025-09-19)* - - - -| Model Name | Config | CPU | GPU A100 (1x) | GPU H100 (1x) | GPU A100 (8x) | GPU H100 (8x) | TPU v2 (8x) | TPU v5e (1x) | -| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | -| **Model** | | | | | | | | | -| [Efficientnet](https://arxiv.org/abs/1905.11946) | ✅ Supported | ✅ Runs | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check |❔ Needs check | ❔ Needs check | +## Model Configuration Support Status +| Model Name | Config Support Status | +| :--- | :--- | +| **Model** | | +| [Efficientnet](https://arxiv.org/abs/1905.11946) | **✅ Supported** | ### Running this model diff --git a/bonsai/models/llada_8b/README.md b/bonsai/models/llada_8b/README.md index b05c5fd7..71a843b5 100644 --- a/bonsai/models/llada_8b/README.md +++ b/bonsai/models/llada_8b/README.md @@ -2,13 +2,12 @@ This directory contains a pure JAX implementation of the [LLaDA diffusion model](https://huggingface.co/GSAI-ML/LLaDA-8B-Instruct), using the [Flax NNX](flax.readthedocs.io/en/stable/index.html) API. -## Tested on: -*(Last Updated: 2025-10-20)* +## Model Configuration Support Status -| Model Name | Config | CPU | GPU A100 (1x) | GPU H100 (1x) | GPU A100 (8x) | GPU H100 (8x) | TPU v2 (8x) | TPU v5e (1x) | -| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | -| **Dense Models** | | | | | | | | | -| [LLaDA-8B-Instruct](https://huggingface.co/GSAI-ML/LLaDA-8B-Instruct) | ✅ Supported | ✅ Runs | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ✅ Runs | +| Model Name | Config Support Status | +| :--- | :--- | +| **Dense Models** | | +| [LLaDA-8B-Instruct](https://huggingface.co/GSAI-ML/LLaDA-8B-Instruct) | **✅ Supported** | ## How to contribute to this model diff --git a/bonsai/models/qwen3/README.md b/bonsai/models/qwen3/README.md index 44c8a474..1699db61 100644 --- a/bonsai/models/qwen3/README.md +++ b/bonsai/models/qwen3/README.md @@ -2,27 +2,21 @@ This directory contains a pure JAX implementation of the [Qwen3 language model](https://qwenlm.github.io/blog/qwen3/), using the [Flax NNX](https://flax.readthedocs.io/en/v0.8.3/experimental/nnx/index.html) API. -> [!IMPORTANT] -> For large-scale high performance use case, please see the [MaxText](https://github.com/AI-Hypercomputer/maxtext?tab=readme-ov-file#getting-started) version. - - -## Tested on: -*(Last Updated: 2025-07-02)* - - - -| Model Name | Config | CPU | GPU A100 (1x) | GPU H100 (1x) | GPU A100 (8x) | GPU H100 (8x) | TPU v2 (8x) | TPU v5e (1x) | -| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | -| **Dense Models** | | | | | | | | | -| [Qwen3-0.6B](https://huggingface.co/Qwen/Qwen3-0.6B) | ✅ Supported | ✅ Runs | ✅ Runs | ✅ Runs | ✅ Runs | ✅ Runs | ✅ Runs | ✅ Runs | -| [Qwen3-1.7B](https://huggingface.co/Qwen/Qwen3-1.7B) | ✅ Supported | ✅ Runs | ✅ Runs | ✅ Runs | ✅ Runs | ✅ Runs | ✅ Runs | ✅ Runs | -| [Qwen3-4B](https://huggingface.co/Qwen/Qwen3-4B) | ✅ Supported | ❔ Needs check | ❔ Needs check | ✅ Runs | ❔ Needs check | ❔ Needs check| ❔ Needs check | ❔ Needs check| -| [Qwen3-8B](https://huggingface.co/Qwen/Qwen3-8B) | ✅ Supported | ❔ Needs check | ❔ Needs check| ✅ Runs | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | -| [Qwen3-14B](https://huggingface.co/Qwen/Qwen3-14B) | ✅ Supported | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | -| [Qwen3-32B](https://huggingface.co/Qwen/Qwen3-32B) | ✅ Supported | ❔ Needs check | ❔ Needs check | ⛔️ Not supported | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | -| **MoE Models** | | | | | | | | | -| [Qwen3-30B-A3B](https://huggingface.co/Qwen/Qwen3-30B-A3B) | 🟡 Not started | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | -| [Qwen3-235B-A22B](https://huggingface.co/Qwen/Qwen3-235B-A22B) | 🟡 Not started | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | + +## Model Configuration Support Status + +| Model Name | Config Support Status | +| :--- | :--- | +| **Dense Models** | | +| [Qwen3-0.6B](https://huggingface.co/Qwen/Qwen3-0.6B) | **✅ Supported** | +| [Qwen3-1.7B](https://huggingface.co/Qwen/Qwen3-1.7B) | **✅ Supported** | +| [Qwen3-4B](https://huggingface.co/Qwen/Qwen3-4B) | **✅ Supported** | +| [Qwen3-8B](https://huggingface.co/Qwen/Qwen3-8B) | **✅ Supported** | +| [Qwen3-14B](https://huggingface.co/Qwen/Qwen3-14B) | **✅ Supported** | +| [Qwen3-32B](https://huggingface.co/Qwen/Qwen3-32B) | **✅ Supported** | +| **MoE Models** | | +| [Qwen3-30B-A3B](https://huggingface.co/Qwen/Qwen3-30B-A3B) | **🟡 Not started** | +| [Qwen3-235B-A22B](https://huggingface.co/Qwen/Qwen3-235B-A22B) | **🟡 Not started** | ### Running this model diff --git a/bonsai/models/resnet/README.md b/bonsai/models/resnet/README.md index 9f299c91..49f0ffb5 100644 --- a/bonsai/models/resnet/README.md +++ b/bonsai/models/resnet/README.md @@ -2,11 +2,11 @@ This directory contains a pure JAX implementation of the [ResNet-50](https://huggingface.co/microsoft/resnet-50) model, using the Flax NNX API. -## Tested on +## Model Configuration Support Status -| Model Name | Config | CPU | GPU A100 (1x) | GPU H100 (1x) | GPU A100 (8x) | GPU H100 (8x) | TPU v2 (8x) | TPU v5e (1x) | -| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | -| [ResNet-50 v1.5](https://huggingface.co/microsoft/resnet-50) | ✅ Supported | ✅ Runs | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | +| Model Name | Config Support Status | +| :--- | :--- | +| [ResNet-50 v1.5](https://huggingface.co/microsoft/resnet-50) | **✅ Supported** | ### Running this model diff --git a/bonsai/models/sam2/README.md b/bonsai/models/sam2/README.md index 2acf5f07..d8b633f4 100644 --- a/bonsai/models/sam2/README.md +++ b/bonsai/models/sam2/README.md @@ -2,15 +2,15 @@ A minimal, readable JAX + Flax NNX re-implementation of Meta’s [Segment Anything 2](https://github.com/facebookresearch/sam2), enabling promptable image and video segmentation. -## Tested on: - -| Model Name | Config | CPU | GPU A100 (1x) | GPU H100 (1x) | GPU A100 (8x) | GPU H100 (8x) | TPU v2 (8x) | TPU v5e (1x) | -|------------------|--------------|:--------:|:-------------:|:-------------:|:-------------:|:-------------:|:-----------:|:------------:| -| **SAM2 Variants**| | | | | | | | | -| `sam2_tiny` | ✅ Supported | ✅ Runs | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | -| `sam2_small` | ✅ Supported | ✅ Runs | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | -| `sam2_baseplus` | ✅ Supported | ✅ Runs | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | -| `sam2_large` | ✅ Supported | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | +## Model Configuration Support Status: SAM2 Variants + +| Model Name | Config Support Status | +| :--- | :--- | +| **SAM2 Variants** | | +| `sam2_tiny` | **✅ Supported** | +| `sam2_small` | **✅ Supported** | +| `sam2_baseplus` | **✅ Supported** | +| `sam2_large` | **✅ Supported** | ### Running this model diff --git a/bonsai/models/unet/README.md b/bonsai/models/unet/README.md index 488b53da..b88ed795 100644 --- a/bonsai/models/unet/README.md +++ b/bonsai/models/unet/README.md @@ -14,15 +14,12 @@ Needs work on: 3. Clean up `UNet_segmentation_example.ipynb` to demonstrate proper parameter loading. +## Model Configuration Support Status -## Tested on: - - - -| Model Name | Config | CPU | GPU A100 (1x) | GPU H100 (1x) | GPU A100 (8x) | GPU H100 (8x) | TPU v2 (8x) | TPU v5e (1x) | -| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | -| **Model** | | | | | | | | | -| [U-Net](https://arxiv.org/abs/1505.04597) | ✅ Supported | ✅ Runs | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check |❔ Needs check | ❔ Needs check | +| Model Name | Config Support Status | +| :--- | :--- | +| **Model** | | +| [U-Net](https://arxiv.org/abs/1505.04597) | **✅ Supported** | ### Running this model diff --git a/bonsai/models/vae/README.md b/bonsai/models/vae/README.md index 73a15209..f6f2d458 100644 --- a/bonsai/models/vae/README.md +++ b/bonsai/models/vae/README.md @@ -3,15 +3,12 @@ This directory contains a pure JAX implementation of the [VAE - Variational Autoencoder](https://arxiv.org/abs/1312.6114), using the [Flax NNX](https://flax.readthedocs.io/en/v0.8.3/experimental/nnx/index.html) API, including procedures to train, save and load weights, and make inference. -## Tested on: -*(Last Updated: 2025-09-19)* +## Model Configuration Support Status - - -| Model Name | Config | CPU | GPU A100 (1x) | GPU H100 (1x) | GPU A100 (8x) | GPU H100 (8x) | TPU v2 (8x) | TPU v5e (1x) | -| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | -| **Model** | | | | | | | | | -| [VAE - Variational Autoencoder](https://arxiv.org/abs/1312.6114) | ✅ Supported | ✅ Runs | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check |❔ Needs check | ❔ Needs check | +| Model Name | Config Support Status | +| :--- | :--- | +| **Model** | | +| [VAE - Variational Autoencoder](https://arxiv.org/abs/1312.6114) | **✅ Supported** | ### Running this model diff --git a/bonsai/models/vgg19/README.md b/bonsai/models/vgg19/README.md index 51a6f370..424cb460 100644 --- a/bonsai/models/vgg19/README.md +++ b/bonsai/models/vgg19/README.md @@ -2,10 +2,11 @@ This directory contains a pure JAX implementation of the [VGG-19](https://huggingface.co/keras/vgg_19_imagenet) model, using the Flax NNX API. -## Tested on -| Model Name | Config | CPU | GPU A100 (1x) | GPU H100 (1x) | GPU A100 (8x) | GPU H100 (8x) | TPU v2 (8x) | TPU v5e (1x) | -|:--------------------------------------------------------| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | -| [VGG-19](https://huggingface.co/keras/vgg_19_imagenet) | ✅ Supported | ✅ Runs | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | +## Model Configuration Support Status + +| Model Name | Config Support Status | +| :--- | :--- | +| [VGG-19](https://huggingface.co/keras/vgg_19_imagenet) | **✅ Supported** | ### Running this model diff --git a/bonsai/models/vit/README.md b/bonsai/models/vit/README.md index 8193f0f5..99451d96 100644 --- a/bonsai/models/vit/README.md +++ b/bonsai/models/vit/README.md @@ -2,11 +2,11 @@ This directory contains a pure JAX implementation of the [ViT](https://huggingface.co/google/vit-base-patch16-224) model, using the Flax NNX API. -## Tested on +## Model Configuration Support Status -| Model Name | Config | CPU | GPU A100 (1x) | GPU H100 (1x) | GPU A100 (8x) | GPU H100 (8x) | TPU v2 (8x) | TPU v5e (1x) | -| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | -| [ViT](https://huggingface.co/google/vit-base-patch16-224) | ✅ Supported | ✅ Runs | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | ❔ Needs check | +| Model Name | Config Support Status | +| :--- | :--- | +| [vit-base-patch16-224](https://huggingface.co/google/vit-base-patch16-224) | **✅ Supported** | ### Running this model