Adding more pipeline retry stats, compress code fixed

Author: jakep-allenai

olmocr/pipeline.py

@@ -271,7 +271,7 @@ async def process_page(args, worker_id: int, pdf_orig_path: str, pdf_local_path:
                 local_image_rotation = page_response.rotation_correction
                 raise ValueError(f"invalid_page rotation for {pdf_orig_path}-{page_num}")
 
-            metrics.add_metrics(completed_pages=1)
+            metrics.add_metrics(**{"completed_pages": 1, f"finished_on_attempt_{attempt}": 1})
             await tracker.track_work(worker_id, f"{pdf_orig_path}-{page_num}", "finished")
             return PageResult(
                 pdf_orig_path,
@@ -1211,6 +1211,19 @@ async def main():
         f"Page Failure rate: {total_metrics.get('failed_pages', 0) / max(total_metrics.get('completed_pages', 0) + total_metrics.get('failed_pages', 0), 1) * 100:.2f}%"
     )
 
+    # Output finished_on_attempt statistics
+    logger.info("\nPages finished by attempt number:")
+    total_finished = sum(total_metrics.get(f'finished_on_attempt_{i}', 0) for i in range(args.max_page_retries))
+    cumulative = 0
+    
+    for i in range(args.max_page_retries):
+        if f'finished_on_attempt_{i}' in total_metrics:
+            count = total_metrics[f'finished_on_attempt_{i}']
+            cumulative += count
+            percentage = (count / total_finished * 100) if total_finished > 0 else 0
+            cumulative_percentage = (cumulative / total_finished * 100) if total_finished > 0 else 0
+            logger.info(f"  Attempt {i}: {count:,} pages ({percentage:.1f}%) - Cumulative: {cumulative:,} ({cumulative_percentage:.1f}%)")
+
     # Output rates
     if "server_input_tokens_per_sec" in rates:
         logger.info(f"Server Input tokens/sec rate: {rates['server_input_tokens_per_sec']:.2f}")

olmocr/train/quantization_configs/qwen2_5vl_w8a8_fp8_kv8.yaml

@@ -0,0 +1,25 @@
+quant_stage:
+    quant_modifiers:
+        QuantizationModifier:
+            ignore: ['re:.*lm_head', 're:model.visual.*']
+            config_groups:
+                group_0:
+                    weights:
+                        num_bits: 8
+                        type: float
+                        strategy: channel
+                        dynamic: false
+                        symmetric: true
+                    input_activations:
+                        num_bits: 8
+                        type: float
+                        strategy: token
+                        dynamic: true
+                        symmetric: true
+                    targets: ["Linear"]
+            kv_cache_scheme:
+                num_bits: 8
+                type: float
+                strategy: tensor
+                dynamic: false
+                symmetric: true

scripts/compress_model.sh

@@ -144,6 +144,10 @@ task_spec_args = {
     "resources": TaskResources(gpu_count=1),
     "constraints": Constraints(cluster=["ai2/ceres-cirrascale", "ai2/jupiter-cirrascale-2"]),
     "result": ResultSpec(path="/noop-results"),
+    "datasets": [
+        DataMount.new(mount_path="/weka/oe-data-default", weka="oe-data-default"),
+        DataMount.new(mount_path="/weka/oe-training-default", weka="oe-training-default"),
+    ]
 }
 
 # Add env vars if AWS credentials exist