Nền tảng OlmoEarth: Suy luận dữ liệu địa không gian ở quy mô hành tinh

Quay lại các bài viết

🌍 Tìm hiểu thêm về Nền tảng OlmoEarth: https://allenai.org/olmoearth

Các mô hình OlmoEarth là dòng mô hình nền tảng quan sát Trái đất của chúng tôi, được huấn luyện trước trên khoảng 10 terabyte dữ liệu vệ tinh đa phương thức. Các chính phủ, tổ chức phi chính phủ và các tổ chức hoạt động vì mục tiêu xã hội khác đã và đang ứng dụng OlmoEarth cho các mục đích như giám sát nạn phá rừng, an ninh lương thực và rủi ro cháy rừng.

Tại Ai2, chúng tôi biết cách huấn luyện và phát hành các mô hình mở mạnh mẽ. Đối với các tổ chức có đội ngũ kỹ thuật vững mạnh, một mô hình mở là tất cả những gì họ cần để bắt đầu. Tuy nhiên, hầu hết các tổ chức trong lĩnh vực môi trường – những đơn vị có vị thế tốt nhất để áp dụng các mô hình này – lại không có cơ sở hạ tầng hoặc đội ngũ kỹ thuật đủ khả năng quản lý toàn bộ vòng đời: từ gắn nhãn dữ liệu, tinh chỉnh mô hình (fine-tuning) đến vận hành suy luận (inference) quy mô lớn. Chúng tôi đã dành hơn một thập kỷ để vận hành các nền tảng như Skylight và EarthRanger, những phần mềm mà người dùng trên khắp thế giới dựa vào mỗi ngày, vì vậy chúng phải hoạt động ổn định liên tục. Kinh nghiệm đó đã dạy chúng tôi rằng để tạo ra tác động cần phải: vận hành các mô hình một cách hiệu quả về chi phí vào đúng thời điểm và địa điểm, giám sát hiệu suất, chuyển đổi dữ liệu thô thành thông tin chi tiết có thể hành động và xác minh rằng các kết quả đầu ra đó thúc đẩy được những mục tiêu mà đối tác mong muốn.

Đó là lý do tại sao chúng tôi xây dựng Nền tảng OlmoEarth: cơ sở hạ tầng để đưa các mô hình không gian địa lý từ giai đoạn tinh chỉnh và đánh giá đến suy luận quy mô lớn.

Suy luận ở quy mô này đặt ra những thách thức riêng. Hình ảnh vệ tinh phải được tìm kiếm và truy cập từ nhiều nhà cung cấp, được căn chỉnh theo các phép chiếu và độ phân giải khác nhau, đồng thời phải được xử lý hiệu quả. Sau đó, các kết quả phải được ghép lại thành các bản đồ nhất quán về mặt địa lý trong khi cơ sở hạ tầng phải tự phục hồi sau những lỗi thường gặp của hệ thống tính toán phân tán.

Hiện nay, nền tảng có thể thực hiện suy luận trên các khu vực quy mô lục địa chỉ trong khoảng một ngày, xử lý hàng chục terabyte hình ảnh với chi phí chỉ bằng một phần nhỏ của một xu trên mỗi km vuông. Việc phát triển nền tảng này đồng nghĩa với việc phải đối mặt với hàng loạt thách thức kỹ thuật mà những người khác làm việc trên các hệ thống không gian địa lý quy mô lớn cũng có khả năng gặp phải. Bài viết này sẽ đi sâu vào những thách thức đó và các giải pháp mà chúng tôi đã đạt được.

Một bản đồ rủi ro cháy rừng gần đây được tạo trên Nền tảng OlmoEarth, kèm theo các số liệu thống kê.

Tại sao suy luận từ dữ liệu vệ tinh lại đầy thách thức

Hầu hết các mô hình ML tiếp nhận vài megabyte dữ liệu và tạo ra kết quả trong chưa đầy một giây—hãy nghĩ đến các LLM xử lý một đoạn văn bản hoặc các mô hình thị giác máy tính phân tích ảnh từ điện thoại thông minh. Suy luận quan sát Trái đất hoạt động ở một quy mô hoàn toàn khác—một tác vụ đơn lẻ tinh chỉnh mô hình nền tảng để đạt hiệu suất tối đa có thể di chuyển hàng terabyte dữ liệu và chạy trong nhiều giờ. Dữ liệu đầu vào có thể bao gồm nhiều dải phổ, loại cảm biến và các mốc thời gian khác nhau trên một khu vực địa lý rộng lớn. Chúng có thể đến từ nhiều nhà cung cấp, mỗi bên sử dụng các phép chiếu và độ phân giải khác nhau, và có thể bao gồm các quan sát bị thiếu hoặc bị che khuất bởi mây. Bản thân kết quả đầu ra là một bản đồ, vì vậy mọi dự đoán phải được căn chỉnh chính xác với cùng một phép chiếu và lưới tọa độ như các khu vực xung quanh nó.

Ngay cả việc thu thập dữ liệu cũng có thể là một thách thức lớn. Các tác vụ dự đoán thường tốn nhiều thời gian tải xuống và chuẩn bị hình ảnh hơn là chạy chính mô hình đó, điều này khiến các đường ống dữ liệu (data pipelines) hiệu quả trở nên cực kỳ quan trọng. Các đường ống đó phải xử lý lưu lượng I/O lớn trong khi vẫn cung cấp khả năng tính toán cần thiết để tái chiếu và lấy mẫu lại hình ảnh.

Phần cứng phù hợp cho tác vụ phù hợp

Vì việc thu thập và chuẩn bị dữ liệu thường chiếm phần lớn thời gian chạy của một tác vụ suy luận, nên việc giao công việc đó cho GPU sẽ khiến phần cứng đắt tiền nhất của hệ thống phải thực hiện các tác vụ phù hợp hơn với CPU. Do đó, chúng tôi chia mỗi tác vụ thành ba giai đoạn, mỗi giai đoạn tương ứng với một cấu hình phần cứng riêng biệt:

Nền tảng OlmoEarth phân phối các giai đoạn này trên nhiều máy trong khi vẫn giữ cho GPU được tận dụng tối đa. Các trình tải dữ liệu đa tiến trình (multiprocess data loaders) liên tục cung cấp dữ liệu cho mỗi GPU, trong khi các kết quả hoàn thiện được truyền trực tiếp đến bộ lưu trữ blob.

Một yêu cầu, hàng trăm worker và hàng nghìn tiến trình

OlmoEarth Run, lớp thực thi của nền tảng dành cho các tác vụ suy luận quy mô lớn, chia khu vực địa lý được bao phủ bởi mỗi tác vụ thành các phân vùng có kích thước phù hợp cho các thực thể tính toán riêng lẻ (worker), sau đó chia nhỏ các phân vùng đó thành các cửa sổ nhỏ hơn mà các mô hình OlmoEarth sẽ xử lý. Vì mỗi cửa sổ có thể được xử lý độc lập trong một lượt truyền tiến (forward pass) riêng biệt, công việc trên một phần của bản đồ không cần phải chờ đợi phần khác.

Trên thực tế, một khu vực có quy mô bằng một tiểu bang có thể trở thành khoảng một trăm phân vùng, trong khi một lượt chạy quy mô lục địa có thể lên tới hàng nghìn. Các phân vùng liền kề chồng lấp nhẹ lên nhau, và chúng tôi điều chỉnh sự chồng lấp đó khi các kết quả đầu ra được lắp ghép để không xuất hiện đường nối nào trong bản đồ raster cuối cùng.

Vì các phân vùng độc lập với nhau, cùng một giai đoạn có thể chạy trên hàng nghìn thực thể tính toán cùng lúc. Gần đây, chúng tôi đã sử dụng phương pháp này để tạo bản đồ rủi ro cháy rừng bao phủ toàn bộ Bắc Mỹ. Ở mức cao điểm, lượt chạy đã sử dụng khoảng 19.600 CPU và 994 GPU song song, với lưu lượng mạng vượt quá 168 GB/s. Mức độ song song đó đã giảm thời gian tính toán tuần tự ước tính từ 4.737 giờ xuống còn khoảng 30,5 giờ thời gian thực – nhanh gấp 155 lần.

Tuy nhiên, khả năng mở rộng (fan-out) không phải là vô hạn. Càng nhiều worker sẽ càng gây áp lực lên hạn ngạch đám mây, vì vậy tính song song là một nút điều chỉnh cho mỗi lượt chạy, một trong số nhiều thông số mà chúng tôi có thể điều chỉnh trên các tác vụ riêng lẻ. Độ phân giải đầu ra đánh đổi giữa khối lượng dữ liệu và khả năng tính toán để lấy độ chi tiết; kích thước mô hình đánh đổi giữa thời gian GPU và độ chính xác; bộ nhớ đệm hình ảnh thô đánh đổi giữa lưu trữ và tốc độ cho các lượt chạy lặp lại trên cùng một khu vực. Cài đặt phù hợp phụ thuộc vào tác vụ – và ngân sách – hiện có.

Tìm kiếm và lấy đúng các pixel cần thiết

Với một khu vực địa lý và phạm vi thời gian nhất định, nền tảng trước tiên phải xác định những cảnh vệ tinh nào nên được đưa vào mô hình. Điều đó có nghĩa là xác định những gì đã được chụp, ở đâu và khi nào trên các nhà cung cấp với các danh mục, định dạng và độ trễ xuất bản khác nhau. Các tiêu chí lựa chọn cũng phụ thuộc vào nguồn: đối với hình ảnh quang học như Sentinel-2, chúng tôi thường ưu tiên các cảnh ít mây nhất hiện có, trong khi đối với radar khẩu độ tổng hợp (SAR), các kênh phân cực khả dụng có thể quan trọng hơn.

Chúng tôi dựa vào các danh mục STAC công khai và các tiêu chuẩn mở bất cứ khi nào có thể. Nhưng một tác vụ suy luận lớn có thể tạo ra hàng nghìn truy vấn siêu dữ liệu cùng lúc – vượt xa khả năng xử lý đồng thời của các dịch vụ bên ngoài như API STAC của ESA hoặc Microsoft Planetary Computer.

Để tránh làm quá tải các dịch vụ đó, Nền tảng OlmoEarth duy trì chỉ mục siêu dữ liệu riêng, được cập nhật khi hình ảnh mới được xuất bản. Đối với các tập dữ liệu được lưu trữ thông qua AWS Open Data, chúng tôi nhận thông báo SNS cho mỗi cảnh mới. Khi một nhà cung cấp không cung cấp luồng thay đổi, chúng tôi sẽ thăm dò chỉ mục thượng nguồn của họ vài phút một lần. Kết quả là, các yêu cầu của chúng tôi gửi đến các dịch vụ bên ngoài tuân theo nhịp độ ổn định của các ấn phẩm mới thay vì sự bùng nổ đột ngột do một tác vụ suy luận lớn tạo ra.

Mỗi mục trong chỉ mục lưu trữ siêu dữ liệu cảnh cùng với các con trỏ đến mọi vị trí nơi các pixel cơ sở khả dụng. Khi chạy, nền tảng chọn nguồn tốt nhất và thực hiện đọc theo cửa sổ (windowed reads) đối với các định dạng tối ưu hóa cho đám mây như COG hoặc Zarr, chỉ truy xuất các byte cần thiết cho một phân vùng nhất định thay vì tải xuống toàn bộ cảnh.

Chỉ mục này cũng hỗ trợ các công cụ chú thích của chúng tôi. Vì nó duy trì các con trỏ đến hình ảnh Sentinel-1, Sentinel-2, Landsat và NISAR ở các định dạng tối ưu hóa cho đám mây, chúng tôi có thể phục vụ các ô (tiles) từ bất kỳ cảnh nào đã được lập chỉ mục thông qua cùng một hệ thống đọc theo cửa sổ mà không cần xây dựng một đường ống tiếp nhận riêng biệt.

Các nhà cung cấp giúp quy trình làm việc này trở nên dễ dàng nhất đều chia sẻ ba đặc điểm mà chúng tôi khuyến nghị là các phương pháp tốt nhất để xuất bản dữ liệu quan sát Trái đất: thông báo dựa trên hàng đợi khi có hình ảnh mới, lưu trữ trên các nền tảng đám mây lớn mà không có giới hạn tốc độ tùy chỉnh hoặc nút thắt cổ chai về tính khả dụng, và các định dạng tối ưu hóa cho đám mây hỗ trợ đọc theo phạm vi (ranged reads).

Xử lý lỗi ở quy mô lớn

Nền tảng OlmoEarth được thiết kế để tự động phục hồi sau các lỗi. Đối với mỗi tác vụ trong một giai đoạn và phân vùng địa lý, nó sẽ tự động cấp phát một máy ảo chạy container Docker runner của chúng tôi. Runner truy xuất các tham số tác vụ, thực thi công việc, trả về kết quả và tắt máy. Vì mọi tác vụ đều có khả năng tái nhập (reentrant) và lũy đẳng (idempotent), các lỗi gián đoạn có thể được xử lý an toàn bằng cách chạy lại tác vụ đó.

Ở quy mô này, lỗi là điều khó tránh khỏi: một nhà cung cấp có thể chậm hoặc tạm thời không khả dụng; siêu dữ liệu có thể chỉ ra rằng hình ảnh tồn tại ngay cả khi một dải hoặc cửa sổ cần thiết bị thiếu; độ che phủ của mây có thể khiến quá ít quan sát khả dụng; hoặc một tác vụ có thể bị treo hoàn toàn. Nền tảng phản hồi bằng cách theo dõi tác vụ, tự động thử lại, chuyển sang nhà cung cấp thay thế khi có thể và phân biệt rõ ràng giữa các lỗi có thể thử lại và lỗi nghiêm trọng. Một quy trình giám sát riêng biệt sẽ phát hiện các runner bị treo hoặc dừng và khởi động lại các tác vụ của chúng.

Hướng đi tiếp theo của chúng tôi

Lộ trình của chúng tôi được định hình bởi những khoảng trống mà các đối tác đã xác định và những khả năng mà họ cần nhất. Trong số các lĩnh vực chúng tôi đang thực hiện:

Đây mới chỉ là sự khởi đầu. Các mô hình nền tảng không gian địa lý, và đặc biệt là việc đưa chúng vào vận hành thực tế, vẫn là một công nghệ mới nổi, và nhiều tổ chức có vị thế tốt nhất để sử dụng chúng – những đơn vị làm việc trong lĩnh vực bảo tồn, an ninh lương thực, ứng phó thảm họa và khí hậu – chưa bao giờ có quyền truy cập vào cơ sở hạ tầng như thế này. Khoảng cách giữa những gì các đội ngũ này cần hiểu về hành tinh và những gì ngân sách cũng như nguồn lực kỹ thuật cho phép họ thực hiện vẫn còn rất lớn.

Chúng tôi đang xây dựng OlmoEarth để giúp thu hẹp khoảng cách đó.