Create moe-in-transformers.py

moe-in-transformers.py

@@ -0,0 +1,316 @@
+"""
+Analyze model introductions in the Transformers repo over the last ~2 years and
+classify each introduced model as "moe" vs "dense" using a heuristic regex.
+
+Outputs (in ./moe_dense_analysis):
+- moe_dense_models_raw.csv            : all models + inferred intro date + moe/dense label
+- moe_dense_models_2y_window.csv      : only models introduced in the last ~2 years
+- moe_dense_2y_timeline.csv           : monthly cumulative counts (moe/dense/total) over the window
+- moe_dense_2y_timeline.png           : plot of cumulative counts
+"""
+
+# -----------------------------------------------------------------------------
+# Imports
+# -----------------------------------------------------------------------------
+import calendar
+import csv
+import datetime as dt
+import re
+import subprocess
+from pathlib import Path
+
+import matplotlib
+matplotlib.use("Agg")  # headless backend for saving figures on CI/servers
+import matplotlib.dates as mdates
+import matplotlib.pyplot as plt
+
+
+# -----------------------------------------------------------------------------
+# Repo paths / output directory
+# -----------------------------------------------------------------------------
+repo = Path(".").resolve()
+
+models_root = repo / "src/transformers/models"
+if not models_root.exists():
+    raise SystemExit("Run this from the transformers repo root.")
+
+out_dir = repo / "moe_dense_analysis"
+out_dir.mkdir(parents=True, exist_ok=True)
+
+
+# -----------------------------------------------------------------------------
+# Date window: last ~2 years from "today"
+# -----------------------------------------------------------------------------
+today = dt.date.today()
+
+# Handle Feb 29 gracefully when subtracting years
+try:
+    start_date = today.replace(year=today.year - 2)
+except ValueError:
+    # If today is Feb 29 and (today.year - 2) is not a leap year, fallback to Feb 28
+    start_date = today.replace(year=today.year - 2, day=28)
+
+end_date = today
+
+
+# -----------------------------------------------------------------------------
+# Discover model directories
+#
+# We consider a directory to be a "model" if it contains modeling_<name>.py
+# (e.g. src/transformers/models/llama/modeling_llama.py)
+# -----------------------------------------------------------------------------
+model_names = []
+for model_dir in sorted(models_root.iterdir()):
+    if not model_dir.is_dir():
+        continue
+    if model_dir.name.startswith("__"):
+        continue
+
+    modeling_file = model_dir / f"modeling_{model_dir.name}.py"
+    if modeling_file.exists():
+        model_names.append(model_dir.name)
+
+model_name_set = set(model_names)
+
+
+# -----------------------------------------------------------------------------
+# Infer intro date per model using git:
+#
+# We use git log restricted to "added files" under src/transformers/models, and
+# record the earliest date where any file under that model directory was added.
+#
+# NOTE: This is a heuristic, not a perfect "model introduced" definition.
+# -----------------------------------------------------------------------------
+git_out = subprocess.run(
+    [
+        "git",
+        "log",
+        "--diff-filter=A",       # only "added file" changes
+        "--name-only",           # list file paths
+        "--format=DATE %ad",     # insert a marker line with the commit date
+        "--date=short",          # YYYY-MM-DD
+        "--",
+        "src/transformers/models",
+    ],
+    cwd=repo,
+    check=True,
+    text=True,
+    capture_output=True,
+).stdout
+
+intro_dates = {}     # model_name -> earliest YYYY-MM-DD date string we observed
+current_date = None  # date string for the current commit chunk in git_out
+
+for raw_line in git_out.splitlines():
+    line = raw_line.strip()
+    if not line:
+        continue
+
+    # Example marker: "DATE 2024-01-10"
+    if line.startswith("DATE "):
+        current_date = line.split(" ", 1)[1]
+        continue
+
+    # Only consider model paths after we've seen a DATE marker
+    if current_date is None:
+        continue
+    if not line.startswith("src/transformers/models/"):
+        continue
+
+    # Expected path structure:
+    # src/transformers/models/<model_name>/...
+    parts = line.split("/")
+    if len(parts) < 4:
+        continue
+
+    model_name = parts[3]
+    if model_name not in model_name_set:
+        continue
+
+    # Keep the earliest date we've seen for this model
+    old = intro_dates.get(model_name)
+    if old is None or current_date < old:
+        intro_dates[model_name] = current_date
+
+
+# -----------------------------------------------------------------------------
+# MoE heuristic:
+#
+# Search for class definitions in modeling_<name>.py where the class name contains
+# MoE/MOE/Moe or Expert/Experts, AND subclasses nn.Module or torch.nn.Module.
+#
+# If we find at least one such class, label model as "moe", else "dense".
+# -----------------------------------------------------------------------------
+moe_class_re = re.compile(
+    r"^class\s+([A-Za-z0-9_]*(?:MoE|MOE|Moe|Expert|Experts)[A-Za-z0-9_]*)"
+    r"\s*\(\s*(?:nn|torch\.nn)\.Module\s*\)\s*:",
+    re.MULTILINE,
+)
+
+records = []
+for model_name in model_names:
+    intro = intro_dates.get(model_name)
+    if intro is None:
+        # If we couldn't find an intro date, skip it (could be missing due to heuristic)
+        continue
+
+    modeling_file = models_root / model_name / f"modeling_{model_name}.py"
+    text = modeling_file.read_text(encoding="utf-8", errors="ignore")
+
+    matches = sorted(set(moe_class_re.findall(text)))
+    label = "moe" if matches else "dense"
+
+    records.append(
+        {
+            "model": model_name,
+            "introduced_date": intro,                 # YYYY-MM-DD (string)
+            "is_moe": label,                          # "moe" or "dense"
+            "moe_class_matches": ";".join(matches),   # matched class names, if any
+            "modeling_file": str(modeling_file.relative_to(repo)),
+        }
+    )
+
+# Sort by intro date then name for stable outputs
+records.sort(key=lambda row: (row["introduced_date"], row["model"]))
+
+
+# -----------------------------------------------------------------------------
+# Restrict to 2-year window
+# -----------------------------------------------------------------------------
+window_records = []
+for row in records:
+    intro_obj = dt.datetime.strptime(row["introduced_date"], "%Y-%m-%d").date()
+    if start_date <= intro_obj <= end_date:
+        row_copy = dict(row)
+        row_copy["intro_obj"] = intro_obj  # store parsed date for comparisons
+        window_records.append(row_copy)
+
+window_records.sort(key=lambda row: (row["intro_obj"], row["model"]))
+
+
+# -----------------------------------------------------------------------------
+# Build monthly timeline points: start_date, then each next month, ending at end_date
+#
+# We try to keep the day-of-month stable (e.g., the 19th of each month), but clamp
+# to the last day of month if needed (e.g., Feb for day=31).
+# -----------------------------------------------------------------------------
+points = [start_date]
+while points[-1] < end_date:
+    last = points[-1]
+
+    # Compute next month safely
+    year = last.year + (last.month // 12)
+    month = 1 if last.month == 12 else last.month + 1
+    day = min(last.day, calendar.monthrange(year, month)[1])
+
+    next_month = dt.date(year, month, day)
+
+    if next_month > end_date:
+        break
+    points.append(next_month)
+
+# Ensure the last point is exactly end_date
+if points[-1] != end_date:
+    points.append(end_date)
+
+
+# -----------------------------------------------------------------------------
+# Compute cumulative counts at each timeline point
+# -----------------------------------------------------------------------------
+timeline_rows = []
+for point in points:
+    moe_cum = sum(
+        1
+        for row in window_records
+        if row["is_moe"] == "moe" and row["intro_obj"] <= point
+    )
+    dense_cum = sum(
+        1
+        for row in window_records
+        if row["is_moe"] == "dense" and row["intro_obj"] <= point
+    )
+
+    timeline_rows.append(
+        {
+            "date": point.isoformat(),
+            "moe_cumulative": moe_cum,
+            "dense_cumulative": dense_cum,
+            "total_cumulative": moe_cum + dense_cum,
+        }
+    )
+
+
+# -----------------------------------------------------------------------------
+# Write CSV outputs
+# -----------------------------------------------------------------------------
+raw_csv = out_dir / "moe_dense_models_raw.csv"
+with raw_csv.open("w", newline="", encoding="utf-8") as f:
+    writer = csv.DictWriter(
+        f,
+        fieldnames=["model", "introduced_date", "is_moe", "moe_class_matches", "modeling_file"],
+    )
+    writer.writeheader()
+    writer.writerows(records)
+
+window_csv = out_dir / "moe_dense_models_2y_window.csv"
+with window_csv.open("w", newline="", encoding="utf-8") as f:
+    writer = csv.DictWriter(
+        f,
+        fieldnames=["model", "introduced_date", "is_moe", "moe_class_matches", "modeling_file"],
+    )
+    writer.writeheader()
+    for row in window_records:
+        copy_row = dict(row)
+        copy_row.pop("intro_obj", None)  # internal-only field
+        writer.writerow(copy_row)
+
+timeline_csv = out_dir / "moe_dense_2y_timeline.csv"
+with timeline_csv.open("w", newline="", encoding="utf-8") as f:
+    writer = csv.DictWriter(
+        f,
+        fieldnames=["date", "moe_cumulative", "dense_cumulative", "total_cumulative"],
+    )
+    writer.writeheader()
+    writer.writerows(timeline_rows)
+
+
+# -----------------------------------------------------------------------------
+# Plot cumulative counts over time
+# -----------------------------------------------------------------------------
+x = [dt.datetime.strptime(row["date"], "%Y-%m-%d").date() for row in timeline_rows]
+# y_dense = [row["dense_cumulative"] for row in timeline_rows]
+y_moe = [row["moe_cumulative"] for row in timeline_rows]
+
+plt.figure(figsize=(11, 6))
+# plt.plot(x, y_dense, label="Dense cumulative", linewidth=2.2)
+plt.plot(x, y_moe, label="MoE cumulative", linewidth=2.2)
+
+# plt.title(f"MoE vs Dense model introductions ({start_date} to {end_date})")
+plt.title(f"MoE model introductions ({start_date} to {end_date})")
+plt.xlabel("Date")
+plt.ylabel("Model count")
+plt.grid(alpha=0.3)
+plt.legend()
+
+ax = plt.gca()
+ax.xaxis.set_major_locator(mdates.MonthLocator(interval=2))
+ax.xaxis.set_major_formatter(mdates.DateFormatter("%Y-%m"))
+plt.xticks(rotation=45, ha="right")
+plt.tight_layout()
+
+plot_png = out_dir / "moe_dense_2y_timeline.png"
+plt.savefig(plot_png, dpi=180)
+
+
+# -----------------------------------------------------------------------------
+# Print summary
+# -----------------------------------------------------------------------------
+dense_total = sum(1 for row in window_records if row["is_moe"] == "dense")
+moe_total = sum(1 for row in window_records if row["is_moe"] == "moe")
+
+print(f"Window: {start_date} -> {end_date}")
+print(f"Introduced in window: dense={dense_total}, moe={moe_total}, total={dense_total + moe_total}")
+print(f"Wrote {raw_csv}")
+print(f"Wrote {window_csv}")
+print(f"Wrote {timeline_csv}")
+print(f"Wrote {plot_png}")