fix: single-device loading progress, warmup status text, yellow unlock animation
- Load model lazily on single device and eval layers one-by-one with progress callback - Separate LOADING from WARMING UP status so warmup shows "Warming up..." not "Loading model into memory..." - Add warm yellow glow effect to unlock animation chips with peak-and-settle highlight Co-Authored-By: Claude Opus 4.6 <[email protected]>
This commit is contained in:
Alex Cheema
co-authored by
Claude Opus 4.6
Claude Opus 4.6
parent
ebf3cdfd72
commit
86e42b8926
@@ -3156,7 +3156,7 @@
|
||||
{onboardingCombinedGB} GB combined
|
||||
</text>
|
||||
|
||||
<!-- Step 2: Models unlocked — simple fade-in -->
|
||||
<!-- Step 2: Models unlocked — staggered slide-up + yellow glow -->
|
||||
{#if unlockedModels.length > 0 && $chipPhase > 0.01}
|
||||
{@const centerX = ($device1X + $device2X) / 2}
|
||||
{@const chipW = 140}
|
||||
@@ -3164,29 +3164,51 @@
|
||||
{@const chipGap = 12}
|
||||
{@const totalW = unlockedModels.length * chipW + (unlockedModels.length - 1) * chipGap}
|
||||
{@const startX = centerX - totalW / 2}
|
||||
<!-- Header -->
|
||||
{@const headerOpacity = Math.min(1, $chipPhase)}
|
||||
<!-- SVG filter for yellow glow -->
|
||||
<defs>
|
||||
<filter id="chip-glow" x="-50%" y="-50%" width="200%" height="200%">
|
||||
<feGaussianBlur in="SourceGraphic" stdDeviation="4" result="blur" />
|
||||
<feColorMatrix in="blur" type="matrix" values="1 0.8 0 0 0 0.8 0.7 0 0 0 0 0 0 0 0 0 0 0 0.4 0" result="glow" />
|
||||
<feMerge>
|
||||
<feMergeNode in="glow" />
|
||||
<feMergeNode in="SourceGraphic" />
|
||||
</feMerge>
|
||||
</filter>
|
||||
</defs>
|
||||
<!-- Header slides up + fades with yellow tint -->
|
||||
{@const headerProgress = Math.min(1, $chipPhase)}
|
||||
{@const headerY = 332 + 12 * (1 - headerProgress)}
|
||||
{@const yellowR = 234}
|
||||
{@const yellowG = 179}
|
||||
{@const yellowB = 8}
|
||||
<text
|
||||
x={centerX}
|
||||
y="332"
|
||||
y={headerY}
|
||||
text-anchor="middle"
|
||||
dominant-baseline="middle"
|
||||
fill="rgba(255,255,255,{0.4 * headerOpacity})"
|
||||
opacity={headerOpacity}
|
||||
fill="rgba({yellowR},{yellowG},{yellowB},{0.5 * headerProgress})"
|
||||
opacity={headerProgress}
|
||||
style="font-size: 10px; font-family: -apple-system, 'SF Pro Display', system-ui, sans-serif; font-weight: 500; letter-spacing: 0.1em;"
|
||||
>
|
||||
NEW MODELS UNLOCKED
|
||||
</text>
|
||||
<!-- Model chips — staggered fade-in -->
|
||||
<!-- Model chips — staggered slide-up + scale + yellow highlight -->
|
||||
{#each unlockedModels as model, i}
|
||||
{@const stagger = i * 0.6}
|
||||
{@const progress = Math.max(0, Math.min(1, $chipPhase - stagger))}
|
||||
{@const modelName = (model.name || model.id.split('/').pop() || '').slice(0, 18)}
|
||||
{@const modelSize = Math.round(getModelSizeGB(model))}
|
||||
{@const slideY = 16 * (1 - progress)}
|
||||
{@const chipScale = 0.85 + 0.15 * progress}
|
||||
<!-- Yellow highlight peaks at ~0.6 progress then settles to subtle -->
|
||||
{@const highlightPeak = progress < 0.6 ? progress / 0.6 : 1 - (progress - 0.6) / 0.4 * 0.6}
|
||||
{@const borderYellow = 0.15 + 0.35 * highlightPeak}
|
||||
{@const fillYellow = 0.02 + 0.06 * highlightPeak}
|
||||
{#if progress > 0}
|
||||
<g
|
||||
transform="translate({startX + i * (chipW + chipGap) + chipW / 2}, 358)"
|
||||
transform="translate({startX + i * (chipW + chipGap) + chipW / 2}, {358 + slideY}) scale({chipScale})"
|
||||
opacity={progress}
|
||||
filter={highlightPeak > 0.3 ? "url(#chip-glow)" : "none"}
|
||||
>
|
||||
<rect
|
||||
x={-chipW / 2}
|
||||
@@ -3194,8 +3216,8 @@
|
||||
width={chipW}
|
||||
height={chipH}
|
||||
rx="15"
|
||||
fill="rgba(255,255,255,0.03)"
|
||||
stroke="rgba(255,255,255,0.12)"
|
||||
fill="rgba({yellowR},{yellowG},{yellowB},{fillYellow})"
|
||||
stroke="rgba({yellowR},{yellowG},{yellowB},{borderYellow})"
|
||||
stroke-width="1"
|
||||
/>
|
||||
<text
|
||||
@@ -3203,7 +3225,7 @@
|
||||
y={modelSize ? -4 : 1}
|
||||
text-anchor="middle"
|
||||
dominant-baseline="middle"
|
||||
fill="rgba(255,255,255,0.8)"
|
||||
fill="rgba(255,255,255,{0.5 + 0.3 * progress})"
|
||||
style="font-size: 10px; font-family: 'SF Mono', ui-monospace, monospace; font-weight: 500;"
|
||||
>
|
||||
{modelName}
|
||||
@@ -3214,7 +3236,7 @@
|
||||
y="8"
|
||||
text-anchor="middle"
|
||||
dominant-baseline="middle"
|
||||
fill="rgba(255,255,255,0.3)"
|
||||
fill="rgba(255,255,255,{0.15 + 0.15 * progress})"
|
||||
style="font-size: 8px; font-family: 'SF Mono', ui-monospace, monospace; font-weight: 400;"
|
||||
>
|
||||
{modelSize} GB
|
||||
@@ -4016,10 +4038,8 @@
|
||||
{@const statusText = downloadInfo.statusText}
|
||||
{@const isDownloading = downloadInfo.isDownloading}
|
||||
{@const isFailed = statusText === "FAILED"}
|
||||
{@const isLoading =
|
||||
statusText === "LOADING" ||
|
||||
statusText === "WARMING UP" ||
|
||||
statusText === "WAITING"}
|
||||
{@const isLoading = statusText === "LOADING"}
|
||||
{@const isWarmingUp = statusText === "WARMING UP" || statusText === "WAITING"}
|
||||
{@const isReady =
|
||||
statusText === "READY" || statusText === "LOADED"}
|
||||
{@const isRunning = statusText === "RUNNING"}
|
||||
@@ -4461,6 +4481,12 @@
|
||||
Loading model into memory...
|
||||
</p>
|
||||
{/if}
|
||||
{:else if isWarmingUp}
|
||||
<p
|
||||
class="text-[11px] text-white/50 leading-relaxed"
|
||||
>
|
||||
Warming up...
|
||||
</p>
|
||||
{:else if isReady || isRunning}
|
||||
<p
|
||||
class="text-[11px] text-green-400/70 leading-relaxed"
|
||||
@@ -4935,10 +4961,8 @@
|
||||
{@const statusText = downloadInfo.statusText}
|
||||
{@const isDownloading = downloadInfo.isDownloading}
|
||||
{@const isFailed = statusText === "FAILED"}
|
||||
{@const isLoading =
|
||||
statusText === "LOADING" ||
|
||||
statusText === "WARMING UP" ||
|
||||
statusText === "WAITING"}
|
||||
{@const isLoading = statusText === "LOADING"}
|
||||
{@const isWarmingUp = statusText === "WARMING UP" || statusText === "WAITING"}
|
||||
{@const isReady =
|
||||
statusText === "READY" || statusText === "LOADED"}
|
||||
{@const isRunning = statusText === "RUNNING"}
|
||||
@@ -5380,6 +5404,12 @@
|
||||
Loading model into memory...
|
||||
</p>
|
||||
{/if}
|
||||
{:else if isWarmingUp}
|
||||
<p
|
||||
class="text-[11px] text-white/50 leading-relaxed"
|
||||
>
|
||||
Warming up...
|
||||
</p>
|
||||
{:else if isReady || isRunning}
|
||||
<p
|
||||
class="text-[11px] text-green-400/70 leading-relaxed"
|
||||
|
||||
@@ -187,7 +187,7 @@ def set_pipeline_prefill(model: nn.Module, is_prefill: bool) -> None:
|
||||
layer.is_prefill = is_prefill
|
||||
|
||||
|
||||
def _inner_model(model: nn.Module) -> nn.Module:
|
||||
def get_inner_model(model: nn.Module) -> nn.Module:
|
||||
inner = getattr(model, "model", None)
|
||||
if isinstance(inner, nn.Module):
|
||||
return inner
|
||||
@@ -205,7 +205,7 @@ def _inner_model(model: nn.Module) -> nn.Module:
|
||||
raise ValueError("Model must either have a 'model' or 'transformer' attribute")
|
||||
|
||||
|
||||
def _get_layers(inner_model_instance: nn.Module) -> list[_LayerCallable]:
|
||||
def get_layers(inner_model_instance: nn.Module) -> list[_LayerCallable]:
|
||||
# Handle both model.layers and model.h cases
|
||||
layers: list[_LayerCallable]
|
||||
if hasattr(inner_model_instance, "layers"):
|
||||
@@ -232,9 +232,9 @@ def pipeline_auto_parallel(
|
||||
Returns:
|
||||
The parallelized model
|
||||
"""
|
||||
inner_model_instance: nn.Module = _inner_model(model)
|
||||
inner_model_instance: nn.Module = get_inner_model(model)
|
||||
|
||||
layers = _get_layers(inner_model_instance)
|
||||
layers = get_layers(inner_model_instance)
|
||||
|
||||
start_layer, end_layer = model_shard_meta.start_layer, model_shard_meta.end_layer
|
||||
device_rank, world_size = model_shard_meta.device_rank, model_shard_meta.world_size
|
||||
@@ -526,7 +526,7 @@ class LlamaShardingStrategy(TensorParallelShardingStrategy):
|
||||
|
||||
|
||||
def _set_layers(model: nn.Module, layers: list[_LayerCallable]) -> None:
|
||||
inner_model_instance = _inner_model(model)
|
||||
inner_model_instance = get_inner_model(model)
|
||||
if hasattr(inner_model_instance, "layers"):
|
||||
inner_model_instance.layers = layers
|
||||
|
||||
|
||||
@@ -58,6 +58,8 @@ from exo.worker.engines.mlx.auto_parallel import (
|
||||
LayerLoadedCallback,
|
||||
TimeoutCallback,
|
||||
eval_with_timeout,
|
||||
get_inner_model,
|
||||
get_layers,
|
||||
pipeline_auto_parallel,
|
||||
tensor_auto_parallel,
|
||||
)
|
||||
@@ -169,6 +171,26 @@ def initialize_mlx(
|
||||
return mlx_distributed_init(bound_instance)
|
||||
|
||||
|
||||
def _eval_layers_with_progress(
|
||||
model: nn.Module,
|
||||
on_layer_loaded: LayerLoadedCallback | None = None,
|
||||
) -> None:
|
||||
"""Evaluate model layers one by one, reporting progress via callback."""
|
||||
try:
|
||||
inner = get_inner_model(model)
|
||||
layers = get_layers(inner)
|
||||
except ValueError:
|
||||
# Model doesn't have standard layer structure — eval all at once
|
||||
mx.eval(model)
|
||||
return
|
||||
|
||||
total = len(layers)
|
||||
for i, layer in enumerate(layers):
|
||||
mx.eval(layer) # type: ignore
|
||||
if on_layer_loaded is not None:
|
||||
on_layer_loaded(i + 1, total)
|
||||
|
||||
|
||||
def load_mlx_items(
|
||||
bound_instance: BoundInstance,
|
||||
group: Group | None,
|
||||
@@ -179,7 +201,9 @@ def load_mlx_items(
|
||||
logger.info(f"Single device used for {bound_instance.instance}")
|
||||
model_path = build_model_path(bound_instance.bound_shard.model_card.model_id)
|
||||
start_time = time.perf_counter()
|
||||
model, _ = load_model(model_path, strict=True)
|
||||
model, _ = load_model(model_path, lazy=True, strict=False)
|
||||
_eval_layers_with_progress(model, on_layer_loaded)
|
||||
mx.eval(model)
|
||||
end_time = time.perf_counter()
|
||||
logger.info(f"Time taken to load model: {(end_time - start_time):.2f}s")
|
||||
tokenizer = get_tokenizer(model_path, bound_instance.bound_shard)
|
||||
|
||||
Reference in New Issue
Block a user